90 lines
1.4 KiB
Go
90 lines
1.4 KiB
Go
package main
|
|
|
|
type TrieNode struct {
|
|
has bool
|
|
successors [26]*TrieNode
|
|
}
|
|
|
|
type Trie struct {
|
|
root TrieNode
|
|
}
|
|
|
|
func NewTrie() Trie {
|
|
return Trie{}
|
|
}
|
|
|
|
// [XXX] Uncomment for LeetCode
|
|
// func Constructor() Trie {
|
|
// return NewTrie()
|
|
// }
|
|
|
|
func (this *Trie) Insert(word string) {
|
|
node := &this.root
|
|
|
|
for _, c := range word {
|
|
idx := int(c) - int('a')
|
|
|
|
if node.successors[idx] == nil {
|
|
node.successors[idx] = new(TrieNode)
|
|
}
|
|
node = node.successors[idx]
|
|
}
|
|
|
|
node.has = true
|
|
}
|
|
|
|
func (this *Trie) findPrefix(prefix string) *TrieNode {
|
|
node := &this.root
|
|
|
|
for _, c := range prefix {
|
|
idx := int(c) - int('a')
|
|
|
|
if node.successors[idx] == nil {
|
|
return nil
|
|
}
|
|
node = node.successors[idx]
|
|
}
|
|
|
|
return node
|
|
}
|
|
|
|
func (this *Trie) Search(word string) bool {
|
|
node := this.findPrefix(word)
|
|
return node != nil && node.has
|
|
}
|
|
|
|
func (this *Trie) StartsWith(prefix string) bool {
|
|
var terminates func(*TrieNode) bool
|
|
terminates = func(node *TrieNode) bool {
|
|
if node == nil {
|
|
return false
|
|
}
|
|
|
|
if node.has {
|
|
return true
|
|
}
|
|
|
|
for _, succ := range node.successors {
|
|
if terminates(succ) {
|
|
return true
|
|
}
|
|
}
|
|
|
|
return false
|
|
}
|
|
|
|
node := this.findPrefix(prefix)
|
|
if node == nil {
|
|
return false
|
|
}
|
|
|
|
return terminates(node)
|
|
}
|
|
|
|
/**
|
|
* Your Trie object will be instantiated and called as such:
|
|
* obj := Constructor();
|
|
* obj.Insert(word);
|
|
* param_2 := obj.Search(word);
|
|
* param_3 := obj.StartsWith(prefix);
|
|
*/
|