Waymaker 与交叉创新中心联合 Logo
← 返回博客列表

字典树（Trie / Prefix Tree）算法原理与 Go 语言实现

ryou •

原文地址：https://www.cnblogs.com/ryuo-ou/p/19814220 同步说明：该文已完整同步到站内博客，便于统一检索和阅读。

引言

在处理大量字符串的集合时，我们经常会遇到两类问题：

  1. 快速查询一个单词是否在集合中存在。
  2. 快速判断集合中是否存在以某个前缀开头的单词。

比如我们在搜索引擎中输入前几个字母，底下的搜索建议是如何快速过滤出来的？拼写检查器如何高效判断单词是否合法？如果使用普通的哈希表（Hash Table），我们可以 \(O(1)\) 时间做精确匹配，但却无法高效完成前缀搜索。

这就引入了本文的主角：字典树（又称 Trie 或 前缀树）。

数据结构分析

字典树的核心思想是空间换时间，利用字符串的公共前缀来降低查询时间的开销，从而达到提高效率的目的。

我们可以这样定义字典树的节点：

type Node struct { children [26]*Node // 指向子节点的指针数组，代表 'a' 到 'z' 26个小写字母 isEnd bool // 标记当前节点是否为一个完整单词的结尾
}

这里我们假设字符集只包含 26 个小写字母。通过字符距离 'a' 的相对偏移量（即 ch - 'a'），我们可以把字符映射到 0-25 的数组索引中。如果在该路径上有后续字符序列，对应的 children[i] 的值就为非空指针。

核心操作

1. 插入（Insert）

插入操作从字典树的根节点开始，遍历待插入字符串的每一个字符：

  • 如果子节点存在，沿着指针向下移动。
  • 如果子节点不存在，创建一个新的子节点，并且向下移动。
  • 遍历结束后，将当前节点的 isEnd 属性设为 true，以标记该路径构成了一个完整的单词。

2. 查找（Search 与 StartsWith）

无论是查找完整单词，还是查找前缀，其核心遍历逻辑是一模一样的。为了复用代码，我们可以抽象出一个 find 方法：

  • 沿着给定的字符串路径在树中向下查找。
  • 只要路径断开（即指向了 nil），说明既没找到单词、也没找到前缀，直接返回 0。
  • 如果成功走完了字符串路径，此时根据当前节点的 isEnd 是否为真：
    • 如果为真：说明不仅该前缀存在，而且树中有以此为结尾的那个完整单词（返回 2）。
    • 如果为假：说明目前走的路径仅仅是别的长单词的前缀（返回 1）。

利用这套精简的返回值，Search 方法只需要判断返回值是否为 2；StartsWith 方法只需判断返回值是否大于 0。

Go 语言实现完整代码

基于上述的思路，我们进行了代码规范化，使其符合 Go 语言惯用形式（Idiomatic Go）：

type Node struct { children [26]*Node isEnd bool
} type Trie struct { root *Node
} func Constructor() Trie { return Trie{root: &Node{}}
} // 插入一个单词到前缀树中
func (t *Trie) Insert(word string) { cur := t.root for _, ch := range word { ch -= 'a' if cur.children[ch] == nil { cur.children[ch] = &Node{} } cur = cur.children[ch] } cur.isEnd = true
} // 通用查找方法定义：
// 返回 0: 既不是单词也不是前缀 (找不到)
// 返回 1: 仅仅是存在的前缀
// 返回 2: 完整匹配存在的单词
func (t *Trie) find(prefix string) int { cur := t.root for _, ch := range prefix { ch -= 'a' if cur.children[ch] == nil { return 0 } cur = cur.children[ch] } if cur.isEnd { return 2 } return 1 } // 查询单词是否存在
func (t *Trie) Search(word string) bool { return t.find(word) == 2
} // 查询前缀是否存在
func (t *Trie) StartsWith(prefix string) bool { return t.find(prefix) > 0
}

复杂度分析

  1. 时间复杂度：无论是 Insert、Search 还是 StartsWith，时间复杂度均为 \(O(M)\)，其中 \(M\) 为该字符串的长度。因为要在树中从上往下进行逐字符匹配，不会发生回溯操作，效率非常高。
  2. 空间复杂度：最坏情况下，我们要插入的所有单词之间都没有公共前缀，此时空间复杂度为 \(O(N \times 26)\)（\(N\) 为所有插入字符串的长度总和）。但实际应用中，由于存在大量公共前缀，实际内存占用要少得多。

总结

字典树巧妙利用了包含关系的公共前缀进行状态转移计算，虽然有着不小的空间内存开销（多叉树指针），但在处理大量字符串前缀匹配时能极大降低查询的时间复杂度。这是一种极为实用且经典的数据结构，在刷算法题和工程实践中都占据很重要的地位。

← 返回首页