u

data-structure/ds-trie.go

@@ -0,0 +1,109 @@
+package ds
+
+// ds-trie.go
+// 前缀树 (Trie),支持任意字符串的插入/查找/前缀匹配/删除。
+// 每个节点 26 个子指针(默认按小写字母 a-z 索引),isEnd 标记单词结尾。
+// 如需支持更多字符集,可改用 map[rune]*TrieNode。
+
+// TrieNode 前缀树节点。
+type TrieNode struct {
+	children [26]*TrieNode
+	isEnd    bool
+}
+
+// Trie 前缀树。默认只支持小写字母 a-z。
+type Trie struct {
+	root *TrieNode
+}
+
+// NewTrie 创建空 Trie。
+func NewTrie() *Trie {
+	return &Trie{root: &TrieNode{}}
+}
+
+// Insert 插入单词。仅支持 a-z,其他字符会 panic(用 rune 索引前先校验)。
+func (t *Trie) Insert(word string) {
+	node := t.root
+	for _, ch := range word {
+		idx := ch - 'a'
+		if idx < 0 || idx >= 26 {
+			panic("trie: 只支持小写字母 a-z,收到非法字符")
+		}
+		if node.children[idx] == nil {
+			node.children[idx] = &TrieNode{}
+		}
+		node = node.children[idx]
+	}
+	node.isEnd = true
+}
+
+// Search 查找完整单词是否存在(必须 isEnd)。
+func (t *Trie) Search(word string) bool {
+	node := t.find(word)
+	return node != nil && node.isEnd
+}
+
+// StartsWith 是否存在以 prefix 为前缀的单词。
+func (t *Trie) StartsWith(prefix string) bool {
+	return t.find(prefix) != nil
+}
+
+// find 沿前缀走到对应节点,不存在返回 nil。
+func (t *Trie) find(s string) *TrieNode {
+	node := t.root
+	for _, ch := range s {
+		idx := ch - 'a'
+		if idx < 0 || idx >= 26 {
+			return nil
+		}
+		node = node.children[idx]
+		if node == nil {
+			return nil
+		}
+	}
+	return node
+}
+
+// Delete 删除单词。仅当单词存在时删除,并回收不再被任何单词共享的节点。
+// 返回是否删除成功(单词存在且被删除)。
+func (t *Trie) Delete(word string) bool {
+	// 先确认单词存在,避免与回收语义混淆。
+	node := t.find(word)
+	if node == nil || !node.isEnd {
+		return false
+	}
+	t.deleteNode(t.root, word, 0)
+	return true
+}
+
+// deleteNode 递归回收:删除 word 沿途不再被需要的节点。
+// 调用前需保证 word 存在。返回值:当前节点是否可被父节点回收。
+func (t *Trie) deleteNode(node *TrieNode, word string, depth int) bool {
+	if node == nil {
+		return false
+	}
+	if depth == len(word) {
+		node.isEnd = false
+		return isEmpty(node)
+	}
+	idx := int(word[depth] - 'a')
+	if idx < 0 || idx >= 26 {
+		return false
+	}
+	child := node.children[idx]
+	if t.deleteNode(child, word, depth+1) {
+		node.children[idx] = nil
+	}
+	// 当前节点可回收:无子节点 且 非单词结尾
+	return isEmpty(node) && !node.isEnd
+}
+
+// isEmpty 节点是否无任何子节点。
+func isEmpty(node *TrieNode) bool {
+	for i := 0; i < 26; i++ {
+		if node.children[i] != nil {
+			return false
+		}
+	}
+	return true
+}