diff --git a/go/search-suggestions-system.go b/go/search-suggestions-system.go new file mode 100644 index 0000000..f9c7e6e --- /dev/null +++ b/go/search-suggestions-system.go @@ -0,0 +1,80 @@ +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) +}