Trie树的Go实现

Trie树是一种多叉树数据结构，常用于存储具有公共前缀的字符串集合，拥有节省空间、快速查找的特性。

受到https://oreki.blog.csdn.net/article/details/109076473的C++实现Trie的启发，尝试使用Go来实现Trie树：

type TrieNode struct {
	data    rune
	isEnd   bool
	subNode map[rune]*TrieNode
}

type TrieTree struct {
	root *TrieNode
}

func NewTrieTree() *TrieTree {
	return &TrieTree{
		root: &TrieNode{
			data:    0,
			isEnd:   false,
			subNode: make(map[rune]*TrieNode),
		},
	}
}

//插入字符串
func (t *TrieTree) Insert(s string) {
	if s == "" {
		return
	}
	cur := t.root
	for _, char := range s {
		if _, ok := cur.subNode[char]; !ok {
			cur.subNode[char] = &TrieNode{
				data:    char,
				isEnd:   false,
				subNode: make(map[rune]*TrieNode),
			}
		}
		cur = cur.subNode[char]
	}
	cur.isEnd = true
}

func (t *TrieTree) Find(s string) bool {
	if s == "" {
		return false
	}
	cur := t.root
	for _, char := range s {
		if _, ok := cur.subNode[char]; !ok {
			return false
		}
		cur = cur.subNode[char]
	}
	//如果要在含hello的Trie树中中找he，当遍历完he的每个字符，都与hello的前缀一致，但是trie树的e的node属性isEnd是false，表明Trie树不含he字符
	return cur.isEnd
}

func (t *TrieTree) StartsWith(prefix string) bool {
	if prefix == "" {
		return false
	}
	cur := t.root
	for _, char := range prefix {
		if _, ok := cur.subNode[char]; !ok {
			return false
		}
		cur = cur.subNode[char]
	}
	return true
}

func (t *TrieTree) GetPrefixWords(prefix string) []string {
	if prefix == "" {
		return nil
	}
	cur := t.root
	for _, char := range prefix {
		if _, ok := cur.subNode[char]; !ok {
			return nil
		}
		cur = cur.subNode[char]
	}

	var result []string
	t.getPrefixWordsHelper(cur, prefix, &result)
	return result
}

func (t *TrieTree) getPrefixWordsHelper(cur *TrieNode, str string, result *[]string) {
	if cur.isEnd {
		*result = append(*result, str)
	}
	for char, subNode := range cur.subNode {
		t.getPrefixWordsHelper(subNode, str+string(char), result)
	}
}

func main() {
	trie := NewTrieTree()

	words := []string{"bat", "batman", "batter", "bath", "apple", "apricot", "banana"}
	for _, word := range words {
		trie.Insert(word)
	}

	fmt.Println(trie.Find("bat"))      // true
	fmt.Println(trie.Find("batter"))   // true
	fmt.Println(trie.Find("batwoman")) // false

	fmt.Println(trie.StartsWith("ba"))  // true
	fmt.Println(trie.StartsWith("bat")) // true
	fmt.Println(trie.StartsWith("foo")) // false

	fmt.Println(trie.GetPrefixWords("ba")) // [banana, bat, batman, batter, bath]
}