实现一个 Trie (前缀树)

Trie.go

// 实现一个 Trie (前缀树)，包含 insert, search, 和 startsWith 这三个操作。

// 示例:

// Trie trie = new Trie();

// trie.insert("apple");
// trie.search("apple");   // 返回 true
// trie.search("app");     // 返回 false
// trie.startsWith("app"); // 返回 true
// trie.insert("app");
// trie.search("app");     // 返回 true
// 说明:

// 你可以假设所有的输入都是由小写字母 a-z 构成的。
// 保证所有输入均为非空字符串。

// 来源：力扣（LeetCode）
// 链接：https://leetcode-cn.com/problems/implement-trie-prefix-tree
// 著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

package main

import (
	"fmt"
)

type Trie struct {
	data     int
	children [26]*Trie
	isEnd    bool
}

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

var aIntVal = int('a')

/** Inserts a word into the trie. */
func (this *Trie) Insert(word string) {
	trie := this

	for _, c := range word {
		idx := int(c) - aIntVal

		if trie.children[idx] == nil {
			trie.children[idx] = &Trie{data: int(c)}
		}
		trie = trie.children[idx]
	}
	trie.isEnd = true
}

/** Returns if the word is in the trie. */
func (this *Trie) Search(word string) bool {
	trie := this.match(word)
	return trie != nil && trie.isEnd 
}

/** Returns if there is any word in the trie that starts with the given prefix. */
func (this *Trie) StartsWith(prefix string) bool {
	trie := this.match(prefix)
	return trie != nil
}

func (this *Trie) match(word string) *Trie {
	trie := this
	for _, c := range word {
		idx := int(c) - aIntVal
		if trie.children[idx]  == nil {
			return nil
		}
		trie = trie.children[idx]
	}
	return trie
}

func main() {
	trie := Constructor()

	trie.Insert("apple")
	fmt.Println(trie.Search("apple"))   // 返回 true
	fmt.Println(trie.Search("app"))     // 返回 false
	fmt.Println(trie.StartsWith("app")) // 返回 true
	trie.Insert("app")
	fmt.Println(trie.Search("app")) // 返回 true

}