LeetCode/go/implement-trie-prefix-tree.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);
 */