LeetCode/go/search-suggestions-system.go

package search_suggestions_system

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) Suggest(word string) [][]string {
	firstThree := func(node *TrieNode, prefix string) []string {
		words := make([]string, 0, 3)

		var dfs func(*TrieNode, string)
		dfs = func(node *TrieNode, prefix string) {
			if len(words) == 3 || node == nil {
				return
			}

			if node.has {
				words = append(words, prefix)
			}

			for i, child := range node.successors {
				if child == nil {
					continue
				}

				dfs(child, prefix+string(rune(i+int('a'))))
			}
		}

		dfs(node, prefix)
		return words
	}

	suggestions := make([][]string, len(word))
	node := &this.root

	for i, c := range word {
		idx := int(c) - int('a')
		if node != nil {
			node = node.successors[idx]
		}

		suggestions[i] = firstThree(node, word[:i+1])
	}

	return suggestions
}

func suggestedProducts(products []string, searchWord string) [][]string {
	trie := Constructor()
	for _, product := range products {
		trie.Insert(product)
	}

	return trie.Suggest(searchWord)
}