leetcode-go/medium/208.implement-trie-prefix-tree.go

type Node struct {
	IsKey bool
	Children *[26]Node
}


type Trie struct {
	Root *Node
}


/** Initialize your data structure here. */
func Constructor() Trie {
	return Trie{&Node{}}
}


/** Inserts a word into the trie. */
func (this *Trie) Insert(word string)  {
	curr := this.Root
	var i int
	for i = range word {
		if curr.Children == nil {
			curr.Children = &[26]Node{}
		}
		curr = &curr.Children[int(word[i]-'a')]
	}
	curr.IsKey = true
}


/** Returns if the word is in the trie. */
func (this *Trie) Search(word string) bool {
	curr := this.Root
	var i int
	for i = range word {
		if curr.Children == nil {
			return false
		}
		curr = &curr.Children[int(word[i]-'a')]
	}
	return curr.IsKey
}


/** Returns if there is any word in the trie that starts with the given prefix. */
func (this *Trie) StartsWith(prefix string) bool {
	curr := this.Root
	for i := range prefix {
		if curr.Children == nil {
			return false
		}
		curr = &curr.Children[int(prefix[i]-'a')]
	}
	return curr.IsKey || curr.Children != nil
}


/**
 * 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);
 */