Trie 树又叫字典树、前缀树、单词查找树,是一种二叉树衍生出来的高级数据结构,主要应用场景是处理字符串前缀相关的操作。
关于Map和Set,是两个抽象数据结构(接口),Map存储一个键值对集合,其中键不重复,Set存储一个不重复的元素集合。
常见的Map和Set的底层实现原理有哈希表和二叉搜索树两种,比如 Java 的 HashMap/HashSet 和 C++ 的 unorderd_map/unordered_set 底层就是用哈希表实现,而 Java 的 TreeMap/TreeSet 和 C++ 的 map/set 底层使用红黑树这种自平衡 BST 实现的。
而本文实现的 TrieSet/TrieMap 底层则用 Trie 树这种结构来实现。
了解数据结构的读者应该知道,本质上Set可以视为一种特殊的Map,Set其实就是Map中的键。
所以本文先实现TrieMap,然后在TrieMap的基础上封装出TrieSet。
Trie 树原理
Trie 树本质上就是一棵从二叉树衍生出来的多叉树。
二叉树节点的代码实现是这样:
/* 基本的二叉树节点 */
class TreeNode {
int val;
TreeNode left, right;
}
其中left, right
存储左右子节点的指针,所以二叉树的结构是这样:
多叉树节点的代码实现是这样:
/* 基本的多叉树节点 */
class TreeNode {
int val;
TreeNode[] children;
}
其中children
数组中存储指向孩子节点的指针,所以多叉树的结构是这样:
而TrieMap
中的树节点TrieNode
的代码实现是这样:
/* Trie 树节点实现 */
class TrieNode<V> {
V val = null;
TrieNode<V>[] children = new TrieNode[256];
}
这个val
字段存储键对应的值,children
数组存储指向子节点的指针。
但是和之前的普通多叉树节点不同,TrieNode
中children
数组的索引是有意义的,代表键中的一个字符。
比如说children[97]
如果非空,说明这里存储了一个字符'a'
,因为'a'
的 ASCII 码为 97。
我们的模板只考虑处理 ASCII 字符,所以children
数组的大小设置为 256
。不过这个可以根据具体问题修改,比如改成更小的数组或者HashMap<Character, TrieNode>
都是一样的效果。
有了以上铺垫,Trie 树的结构是这样的:
一个节点有 256 个子节点指针,但大多数时候都是空的,可以省略掉不画,所以一般你看到的 Trie 树长这样:
这是在TrieMap<Integer>
中插入一些键值对后的样子,白色节点代表val
字段为空,橙色节点代表val
字段非空。
这里要特别注意,TrieNode
节点本身只存储val
字段,并没有一个字段来存储字符,字符是通过子节点在父节点的children
数组中的索引确定的。
形象理解就是,Trie
树用「树枝」存储字符串(键),用「节点」存储字符串(键)对应的数据(值)。所以我在图中把字符标在树枝,键对应的值val
标在节点上:
TrieMap/TrieSet API 及实现
API
// 底层用 Trie 树实现的键值映射
// 键为 String 类型,值为类型 V
class TrieMap<V> {
/***** 增/改 *****/
// 在 Map 中添加 key
public void put(String key, V val);
/***** 删 *****/
// 删除键 key 以及对应的值
public void remove(String key);
/***** 查 *****/
// 搜索 key 对应的值,不存在则返回 null
// get("the") -> 4
// get("tha") -> null
public V get(String key);
// 判断 key 是否存在在 Map 中
// containsKey("tea") -> false
// containsKey("team") -> true
public boolean containsKey(String key);
// 在 Map 的所有键中搜索 query 的最短前缀
// shortestPrefixOf("themxyz") -> "the"
public String shortestPrefixOf(String query);
// 在 Map 的所有键中搜索 query 的最长前缀
// longestPrefixOf("themxyz") -> "them"
public String longestPrefixOf(String query);
// 搜索所有前缀为 prefix 的键
// keysWithPrefix("th") -> ["that", "the", "them"]
public List<String> keysWithPrefix(String prefix);
// 判断是和否存在前缀为 prefix 的键
// hasKeyWithPrefix("tha") -> true
// hasKeyWithPrefix("apple") -> false
public boolean hasKeyWithPrefix(String prefix);
// 通配符 . 匹配任意字符,搜索所有匹配的键
// keysWithPattern("t.a.") -> ["team", "that"]
public List<String> keysWithPattern(String pattern);
// 通配符 . 匹配任意字符,判断是否存在匹配的键
// hasKeyWithPattern(".ip") -> true
// hasKeyWithPattern(".i") -> false
public boolean hasKeyWithPattern(String pattern);
// 返回 Map 中键值对的数量
public int size();
}
实现:
class TrieMap<V> {
// ASCII 码个数
private static final int R = 256;
// 当前存在 Map 中的键值对个数
private int size = 0;
// Trie 树的根节点
private TrieNode<V> root = null;
private static class TrieNode<V> {
V val = null;
TrieNode<V>[] children = new TrieNode[R];
}
/***** 增/改 *****/
// 在 map 中添加或修改键值对
public void put(String key, V val) {
if (!containsKey(key)) {
// 新增键值对
size++;
}
// 需要一个额外的辅助函数,并接收其返回值
root = put(root, key, val, 0);
}
// 定义:向以 node 为根的 Trie 树中插入 key[i..],返回插入完成后的根节点
private TrieNode<V> put(TrieNode<V> node, String key, V val, int i) {
if (node == null) {
// 如果树枝不存在,新建
node = new TrieNode<>();
}
if (i == key.length()) {
// key 的路径已插入完成,将值 val 存入节点
node.val = val;
return node;
}
char c = key.charAt(i);
// 递归插入子节点,并接收返回值
node.children[c] = put(node.children[c], key, val, i + 1);
return node;
}
/***** 删 *****/
// 在 Map 中删除 key
public void remove(String key) {
if (!containsKey(key)) {
return;
}
// 递归修改数据结构要接收函数的返回值
root = remove(root, key, 0);
size--;
}
// 定义:在以 node 为根的 Trie 树中删除 key[i..],返回删除后的根节点
private TrieNode<V> remove(TrieNode<V> node, String key, int i) {
if (node == null) {
return null;
}
if (i == key.length()) {
// 找到了 key 对应的 TrieNode,删除 val
node.val = null;
} else {
char c = key.charAt(i);
// 递归去子树进行删除
node.children[c] = remove(node.children[c], key, i + 1);
}
// 后序位置,递归路径上的节点可能需要被清理
if (node.val != null) {
// 如果该 TireNode 存储着 val,不需要被清理
return node;
}
// 检查该 TrieNode 是否还有后缀
for (int c = 0; c < R; c++) {
if (node.children[c] != null) {
// 只要存在一个子节点(后缀树枝),就不需要被清理
return node;
}
}
// 既没有存储 val,也没有后缀树枝,则该节点需要被清理
return null;
}
/***** 查 *****/
// 搜索 key 对应的值,不存在则返回 null
public V get(String key) {
// 从 root 开始搜索 key
TrieNode<V> x = getNode(root, key);
if (x == null || x.val == null) {
// x 为空或 x 的 val 字段为空都说明 key 没有对应的值
return null;
}
return x.val;
}
// 判断 key 是否存在在 Map 中
public boolean containsKey(String key) {
return get(key) != null;
}
// 判断是和否存在前缀为 prefix 的键
public boolean hasKeyWithPrefix(String prefix) {
// 只要能找到一个节点,就是存在前缀
return getNode(root, prefix) != null;
}
// 在所有键中寻找 query 的最短前缀
public String shortestPrefixOf(String query) {
TrieNode<V> p = root;
// 从节点 node 开始搜索 key
for (int i = 0; i < query.length(); i++) {
if (p == null) {
// 无法向下搜索
return "";
}
if (p.val != null) {
// 找到一个键是 query 的前缀
return query.substring(0, i);
}
// 向下搜索
char c = query.charAt(i);
p = p.children[c];
}
if (p != null && p.val != null) {
// 如果 query 本身就是一个键
return query;
}
return "";
}
// 在所有键中寻找 query 的最长前缀
public String longestPrefixOf(String query) {
TrieNode<V> p = root;
// 记录前缀的最大长度
int max_len = 0;
// 从节点 node 开始搜索 key
for (int i = 0; i < query.length(); i++) {
if (p == null) {
// 无法向下搜索
break;
}
if (p.val != null) {
// 找到一个键是 query 的前缀,更新前缀的最大长度
max_len = i;
}
// 向下搜索
char c = query.charAt(i);
p = p.children[c];
}
if (p != null && p.val != null) {
// 如果 query 本身就是一个键
return query;
}
return query.substring(0, max_len);
}
// 搜索前缀为 prefix 的所有键
public List<String> keysWithPrefix(String prefix) {
List<String> res = new LinkedList<>();
// 找到匹配 prefix 在 Trie 树中的那个节点
TrieNode<V> x = getNode(root, prefix);
if (x == null) {
return res;
}
// DFS 遍历以 x 为根的这棵 Trie 树
traverse(x, new StringBuilder(prefix), res);
return res;
}
// 遍历以 node 节点为根的 Trie 树,找到所有键
private void traverse(TrieNode<V> node, StringBuilder path, List<String> res) {
if (node == null) {
// 到达 Trie 树底部叶子结点
return;
}
if (node.val != null) {
// 找到一个 key,添加到结果列表中
res.add(path.toString());
}
// 回溯算法遍历框架
for (char c = 0; c < R; c++) {
// 做选择
path.append(c);
traverse(node.children[c], path, res);
// 撤销选择
path.deleteCharAt(path.length() - 1);
}
}
// 通配符 . 匹配任意字符
public List<String> keysWithPattern(String pattern) {
List<String> res = new LinkedList<>();
traverse(root, new StringBuilder(), pattern, 0, res);
return res;
}
// 遍历函数,尝试在「以 node 为根的 Trie 树中」匹配 pattern[i..]
private void traverse(TrieNode<V> node, StringBuilder path, String pattern, int i, List<String> res) {
if (node == null) {
// 树枝不存在,即匹配失败
return;
}
if (i == pattern.length()) {
// pattern 匹配完成
if (node.val != null) {
// 如果这个节点存储着 val,则找到一个匹配的键
res.add(path.toString());
}
return;
}
char c = pattern.charAt(i);
if (c == '.') {
// pattern[i] 是通配符,可以变化成任意字符
// 多叉树(回溯算法)遍历框架
for (char j = 0; j < R; j++) {
path.append(j);
traverse(node.children[j], path, pattern, i + 1, res);
path.deleteCharAt(path.length() - 1);
}
} else {
// pattern[i] 是普通字符 c
path.append(c);
traverse(node.children[c], path, pattern, i + 1, res);
path.deleteCharAt(path.length() - 1);
}
}
// 判断是和否存在前缀为 prefix 的键
public boolean hasKeyWithPattern(String pattern) {
// 从 root 节点开始匹配 pattern[0..]
return hasKeyWithPattern(root, pattern, 0);
}
// 函数定义:从 node 节点开始匹配 pattern[i..],返回是否成功匹配
private boolean hasKeyWithPattern(TrieNode<V> node, String pattern, int i) {
if (node == null) {
// 树枝不存在,即匹配失败
return false;
}
if (i == pattern.length()) {
// 模式串走到头了,看看匹配到的是否是一个键
return node.val != null;
}
char c = pattern.charAt(i);
// 没有遇到通配符
if (c != '.') {
// 从 node.children[c] 节点开始匹配 pattern[i+1..]
return hasKeyWithPattern(node.children[c], pattern, i + 1);
}
// 遇到通配符
for (int j = 0; j < R; j++) {
// pattern[i] 可以变化成任意字符,尝试所有可能,只要遇到一个匹配成功就返回
if (hasKeyWithPattern(node.children[j], pattern, i + 1)) {
return true;
}
}
// 都没有匹配
return false;
}
// 从节点 node 开始搜索 key,如果存在返回对应节点,否则返回 null
private TrieNode<V> getNode(TrieNode<V> node, String key) {
TrieNode<V> p = node;
// 从节点 node 开始搜索 key
for (int i = 0; i < key.length(); i++) {
if (p == null) {
// 无法向下搜索
return null;
}
// 向下搜索
char c = key.charAt(i);
p = p.children[c];
}
return p;
}
public int size() {
return size;
}
}
Set实现
class TrieSet {
// 底层用一个 TrieMap,键就是 TrieSet,值仅仅起到占位的作用
// 值的类型可以随便设置,我参考 Java 标准库设置成 Object
private final TrieMap<Object> map = new TrieMap<>();
/***** 增 *****/
// 在集合中添加元素 key
public void add(String key) {
map.put(key, new Object());
}
/***** 删 *****/
// 从集合中删除元素 key
public void remove(String key) {
map.remove(key);
}
/***** 查 *****/
// 判断元素 key 是否存在集合中
public boolean contains(String key) {
return map.containsKey(key);
}
// 在集合中寻找 query 的最短前缀
public String shortestPrefixOf(String query) {
return map.shortestPrefixOf(query);
}
// 在集合中寻找 query 的最长前缀
public String longestPrefixOf(String query) {
return map.longestPrefixOf(query);
}
// 在集合中搜索前缀为 prefix 的所有元素
public List<String> keysWithPrefix(String prefix) {
return map.keysWithPrefix(prefix);
}
// 判断集合中是否存在前缀为 prefix 的元素
public boolean hasKeyWithPrefix(String prefix) {
return map.hasKeyWithPrefix(prefix);
}
// 通配符 . 匹配任意字符,返回集合中匹配 pattern 的所有元素
public List<String> keysWithPattern(String pattern) {
return map.keysWithPattern(pattern);
}
// 通配符 . 匹配任意字符,判断集合中是否存在匹配 pattern 的元素
public boolean hasKeyWithPattern(String pattern) {
return map.hasKeyWithPattern(pattern);
}
// 返回集合中元素的个数
public int size() {
return map.size();
}
}