105 lines
1.8 KiB
Go
105 lines
1.8 KiB
Go
package q208
|
|
|
|
type tNode struct {
|
|
str string
|
|
childs [26]*tNode
|
|
hasWord bool
|
|
}
|
|
|
|
func (n *tNode) index() int { return index(n.str) }
|
|
|
|
func (n *tNode) split(prefixLen int) (*tNode, *tNode) {
|
|
if len(n.str) < prefixLen {
|
|
panic("cannot split node")
|
|
}
|
|
if len(n.str) == prefixLen {
|
|
return n, nil
|
|
}
|
|
|
|
prefix := n.str[:prefixLen]
|
|
n.str = n.str[prefixLen:]
|
|
p := &tNode{str: prefix}
|
|
p.childs[n.index()] = n
|
|
return p, n
|
|
}
|
|
|
|
func commonPrefix(a, b string) int {
|
|
minLen := min(len(a), len(b))
|
|
for i := range minLen {
|
|
if a[i] != b[i] {
|
|
return i
|
|
}
|
|
}
|
|
return minLen
|
|
}
|
|
|
|
func index(str string) int {
|
|
if len(str) > 0 {
|
|
return int(str[0] - 'a')
|
|
}
|
|
return -1 // root
|
|
}
|
|
|
|
type Trie struct {
|
|
root *tNode
|
|
}
|
|
|
|
func Constructor() Trie {
|
|
return Trie{
|
|
root: &tNode{},
|
|
}
|
|
}
|
|
|
|
func (t *Trie) Insert(word string) {
|
|
node := t.root
|
|
for len(word) > 0 {
|
|
idx := index(word)
|
|
next := node.childs[idx]
|
|
if next != nil {
|
|
pflen := commonPrefix(word, next.str)
|
|
next, _ = next.split(pflen)
|
|
node.childs[next.index()] = next
|
|
word = word[pflen:]
|
|
} else {
|
|
next = &tNode{str: word}
|
|
word = ""
|
|
node.childs[idx] = next
|
|
}
|
|
node = next
|
|
}
|
|
node.hasWord = true
|
|
}
|
|
|
|
func (t *Trie) Search(word string) bool {
|
|
node := t.root
|
|
for len(word) > 0 {
|
|
next := node.childs[index(word)]
|
|
if next == nil {
|
|
return false
|
|
}
|
|
pflen := commonPrefix(word, next.str)
|
|
if pflen != len(next.str) {
|
|
return false
|
|
}
|
|
word = word[pflen:]
|
|
node = next
|
|
}
|
|
return node.hasWord
|
|
}
|
|
|
|
func (t *Trie) StartsWith(prefix string) bool {
|
|
node := t.root
|
|
for len(prefix) > 0 {
|
|
next := node.childs[index(prefix)]
|
|
if next == nil {
|
|
return false
|
|
}
|
|
pflen := commonPrefix(prefix, next.str)
|
|
if pflen != len(next.str) && pflen != len(prefix) {
|
|
return false
|
|
}
|
|
prefix = prefix[pflen:]
|
|
node = next
|
|
}
|
|
return true
|
|
}
|