110 lines
2.7 KiB
Go
110 lines
2.7 KiB
Go
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
|
|
}
|