support customized mask char on trie

This commit is contained in:
kevin
2020-08-19 14:54:59 +08:00
parent 2673dbc6e1
commit 2ebb5b6b58
2 changed files with 26 additions and 8 deletions

View File

@@ -2,7 +2,11 @@ package stringx
import "github.com/tal-tech/go-zero/core/lang"
const defaultMask = '*'
type (
TrieOption func(trie *trieNode)
Trie interface {
Filter(text string) (string, []string, bool)
FindKeywords(text string) []string
@@ -10,6 +14,7 @@ type (
trieNode struct {
node
mask rune
}
scope struct {
@@ -18,8 +23,15 @@ type (
}
)
func NewTrie(words []string) Trie {
func NewTrie(words []string, opts ...TrieOption) Trie {
n := new(trieNode)
for _, opt := range opts {
opt(n)
}
if n.mask == 0 {
n.mask = defaultMask
}
for _, word := range words {
n.add(word)
}
@@ -114,6 +126,12 @@ func (n *trieNode) findKeywordScopes(chars []rune) []scope {
func (n *trieNode) replaceWithAsterisk(chars []rune, start, stop int) {
for i := start; i < stop; i++ {
chars[i] = '*'
chars[i] = n.mask
}
}
func WithMask(mask rune) TrieOption {
return func(n *trieNode) {
n.mask = mask
}
}