LeetCode/go/search-suggestions-system.go

package main

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 := NewTrie()
	for _, product := range products {
		trie.Insert(product)
	}

	return trie.Suggest(searchWord)
}