Trie树
一句话说明
Trie 的本质是“把字符串拆成一条条字符路径”,让公共前缀只存一份。
为什么它适合做前缀查询
假设插入:
catcarcarddog
结构会像这样:
root
├─ c
│ └─ a
│ ├─ t
│ └─ r
│ └─ d
└─ d
└─ o
└─ g公共前缀 ca 只存一次,所以:
- 查单词:沿字符一路往下走
- 查前缀:同样一路往下走,只是不用要求最后必须是结尾
Go 代码:基础模板
这里用 map[rune]*TrieNode,字符集更灵活;如果题目明确只有小写字母,再换成长度 26 的数组会更快。
type TrieNode struct {
Children map[rune]*TrieNode
IsEnd bool
}
type Trie struct {
Root *TrieNode
}
func NewTrie() *Trie {
return &Trie{
Root: &TrieNode{Children: make(map[rune]*TrieNode)},
}
}
func (t *Trie) Insert(word string) {
node := t.Root
for _, ch := range word {
if node.Children[ch] == nil {
node.Children[ch] = &TrieNode{Children: make(map[rune]*TrieNode)}
}
node = node.Children[ch]
}
node.IsEnd = true
}
func (t *Trie) Search(word string) bool {
node := t.Root
for _, ch := range word {
if node.Children[ch] == nil {
return false
}
node = node.Children[ch]
}
return node.IsEnd
}
func (t *Trie) StartsWith(prefix string) bool {
node := t.Root
for _, ch := range prefix {
if node.Children[ch] == nil {
return false
}
node = node.Children[ch]
}
return true
}如果只有小写字母
type ArrayTrieNode struct {
Children [26]*ArrayTrieNode
IsEnd bool
}
type ArrayTrie struct {
Root *ArrayTrieNode
}
func NewArrayTrie() *ArrayTrie {
return &ArrayTrie{Root: &ArrayTrieNode{}}
}
func (t *ArrayTrie) Insert(word string) {
node := t.Root
for i := 0; i < len(word); i++ {
idx := word[i] - 'a'
if node.Children[idx] == nil {
node.Children[idx] = &ArrayTrieNode{}
}
node = node.Children[idx]
}
node.IsEnd = true
}这种实现的好处是:
- 常数更小
- 不需要哈希
代价是:
- 只能处理固定字符集
- 空间更“硬”
高级功能:删除单词
删除不是简单地把路径全删掉,因为别的单词可能还在共用前缀。
func (t *Trie) Delete(word string) bool {
var dfs func(node *TrieNode, chars []rune, index int) bool
dfs = func(node *TrieNode, chars []rune, index int) bool {
if index == len(chars) {
if !node.IsEnd {
return false
}
node.IsEnd = false
return len(node.Children) == 0
}
child := node.Children[chars[index]]
if child == nil {
return false
}
shouldDelete := dfs(child, chars, index+1)
if shouldDelete {
delete(node.Children, chars[index])
}
return !node.IsEnd && len(node.Children) == 0
}
chars := []rune(word)
dfs(t.Root, chars, 0)
return true
}获取所有前缀匹配的单词
func (t *Trie) WordsWithPrefix(prefix string) []string {
node := t.Root
for _, ch := range prefix {
if node.Children[ch] == nil {
return nil
}
node = node.Children[ch]
}
result := make([]string, 0)
path := []rune(prefix)
var dfs func(cur *TrieNode)
dfs = func(cur *TrieNode) {
if cur.IsEnd {
result = append(result, string(path))
}
for ch, child := range cur.Children {
path = append(path, ch)
dfs(child)
path = path[:len(path)-1]
}
}
dfs(node)
return result
}经典应用
LeetCode 211:支持 . 通配符
type WordDictionary struct {
trie *Trie
}
func NewWordDictionary() *WordDictionary {
return &WordDictionary{trie: NewTrie()}
}
func (wd *WordDictionary) AddWord(word string) {
wd.trie.Insert(word)
}
func (wd *WordDictionary) Search(word string) bool {
chars := []rune(word)
var dfs func(node *TrieNode, index int) bool
dfs = func(node *TrieNode, index int) bool {
if index == len(chars) {
return node.IsEnd
}
ch := chars[index]
if ch == '.' {
for _, child := range node.Children {
if dfs(child, index+1) {
return true
}
}
return false
}
child := node.Children[ch]
if child == nil {
return false
}
return dfs(child, index+1)
}
return dfs(wd.trie.Root, 0)
}LeetCode 648:替换单词
func replaceWords(dictionary []string, sentence string) string {
trie := NewTrie()
for _, word := range dictionary {
trie.Insert(word)
}
words := strings.Split(sentence, " ")
for i, word := range words {
node := trie.Root
prefix := make([]rune, 0)
replaced := false
for _, ch := range word {
if node.Children[ch] == nil {
break
}
node = node.Children[ch]
prefix = append(prefix, ch)
if node.IsEnd {
words[i] = string(prefix)
replaced = true
break
}
}
if !replaced {
words[i] = word
}
}
return strings.Join(words, " ")
}Trie 和哈希表怎么选
哈希表更适合:
- 精确查单词
Trie 更适合:
- 前缀查询
- 字典序遍历
- 自动补全
- 模糊匹配
易错点
search和startsWith的区别只在最后是否要求IsEnd = true。- 删除单词时不能把公共前缀误删掉。
- 如果题目字符集固定,优先考虑数组版节点。
- Go 里
range string按rune走,和按字节走要分清。