diff --git a/go/implement-trie-prefix-tree.go b/go/implement-trie-prefix-tree.go new file mode 100644 index 0000000..9e36542 --- /dev/null +++ b/go/implement-trie-prefix-tree.go @@ -0,0 +1,85 @@ +package implement_trie_prefix_tree + +type TrieNode struct { + has bool + successors [26]*TrieNode +} + +type Trie struct { + root TrieNode +} + +func Constructor() Trie { + return Trie{} +} + +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); + */