dcg/pkg/trie/trie.go

package trie

type (
	// trieNode 节点
	trieNode struct {
		char     string             // unicode
		isEnding bool               // 是否单词结尾
		children map[rune]*trieNode // 子节点
	}

	// Trie 字典树(前缀树)(多叉树)
	Trie struct {
		root *trieNode // 根节点指针
	}
)

func NewTrieNode(char string) *trieNode {
	return &trieNode{
		char:     char,
		isEnding: false,
		children: make(map[rune]*trieNode),
	}
}

func NewTrie() *Trie {
	// 初始化根
	trieNode := NewTrieNode("/")
	return &Trie{trieNode}
}

// Insert 插入一个单词
func (t *Trie) Insert(word string) {
	node := t.root
	for _, code := range word {
		v, ok := node.children[code]
		if !ok {
			v = NewTrieNode(string(code))
			node.children[code] = v
		}
		node = v
	}
	node.isEnding = true
}

func (t *Trie) Exists(word string) bool {
	node := t.root
	for _, code := range word {
		v, ok := node.children[code]
		if !ok {
			return false
		}
		node = v
	}
	if node.isEnding == false {
		return false
	}
	return true
}