数据结构之Trie

封面画师：ツチヤ     封面ID：79601805

本文参考视频：小马哥教育(SEEMYGO) 2019年 恋上数据结构与算法（第一季）

源码仓库：mofan212/data-structure-and-algorithm (github.com)

辅助学习网址：数据结构和算法动态可视化

0. 需求分析

现在有一需求 ：判断一堆不重复的字符串是否以某个前缀开头？

由于字符串需要是不重复的，因此可以用Set \ Map存储字符串，然后需要遍历所有字符串，并对它们进行判断，这样的话时间复杂度与元素数量（字符串）的个数相关，为O(n)。

那有没有更优的数据结构实现前缀搜索呢？😑

那肯定有啊，天降猛男——Trie登场！👊

1. Trie

基本概念

Trie（\traɪ\），字典树，又称前缀树（Prefix Tree）、单词查找树，注意和Tree（\triː\）发音的区别。

Trie 搜索字符串的效率主要跟字符串的长度有关 。我们最开始提出的方法的搜索效率和元素数量有关，就是说如果元素有100w个，那么就有可能需要遍历100w个元素。两者相比，单看文字描述就觉得Trie的效率要高很多。

自身结构

那么 Trie 是怎么做到的呢？

先得说明一点，Trie 是一种树形结构，但是它不是二叉树，它是多叉树。

假设使用 Trie 存储 cat、dog、doggy、does、cast、add 六个单词，最终得到的 Trie 如下：

针对上面给出的 Trie，我们可以认为根节点没有存储任何元素，而其每个子节点都存储了一个字母。

假设这个时候我们需要查找单词【dog】，从根节点出发，先搜索字母 d，找到后再搜索 o ，最后搜索 g ，这样一来就可以找到单词【dog】了。真是因为这种查找方式，使得Trie 搜索字符串的效率主要跟字符串的长度有关 。

2. 编码实现

2.1 接口设计

以下接口仅供参考，在此提供两套接口。

第一套接口如下：

int size();
boolean isEmpty();
void clear();
boolean contains(String str);
void add(String str);
void remove(String str);
// 判断前缀字符串是否为prefix
boolean startsWith(String prefix); 

第二套接口如下：

int size();
boolean isEmpty();
void clear();
boolean contains(String str);
V add(String str, V value);
V remove(String str);
// 判断前缀字符串是否为prefix
boolean startsWith(String prefix);

两套接口的区别在于添加和删除方法。

对于第一套接口：添加方法传递的参数只有一个，且类型为String，删除方法返回值类型为void。

对于第二套接口：添加方法多传递了一个参数value，类型为范型V，删除方法的返回值类型为V。

这两套接口就类似与 Set 和 Map ，具体使用哪套视情况而定。

---

我们在进行编码时选择第二套进行编码，这样的话，第一套的编码就很简单了。

接口如下：

/**
 * @author 默烦
 * @date 2020/8/7
 */
public class Trie<V> {

    public int size() {
        return 0;
    }

    public boolean isEmpty() {
        return false;
    }

    public void clear() {

    }

    public V get(String key) {
        return null;
    }

    public boolean contains(String key) {
        return false;
    }

    public V add(String key, V value) {
        return null;
    }

    public V remove(String key) {
        return null;
    }

    public boolean startsWith(String prefix) {
        return false;
    }
}

我们再写一下测试类，后续在对上述接口进行实现：

static void test(){
    Trie<Integer> trie = new Trie<>();
    trie.add("cat", 1);
    trie.add("dog", 2);
    trie.add("catalog", 3);
    trie.add("cast", 4);
    trie.add("默烦", 5);
    System.out.println(trie.size() == 5);
    System.out.println(trie.startsWith("do"));
    System.out.println(trie.startsWith("c"));
    System.out.println(trie.startsWith("ca"));
    System.out.println(trie.startsWith("cat"));
    System.out.println(trie.startsWith("cata"));
    System.out.println(!trie.startsWith("hehe"));
    System.out.println(trie.get("默烦") == 5);
    System.out.println(trie.remove("cat") == 1);
    System.out.println(trie.remove("catalog") == 3);
    System.out.println(trie.remove("cast") == 4);
    System.out.println(trie.size() == 2);
    System.out.println(trie.startsWith("默"));
    System.out.println(trie.startsWith("do"));
    System.out.println(!trie.startsWith("c"));
}

2.2 Node设计

假设使用 Trie 存储 cat、dog、doggy、does、cast、add 六个单词，最终得到的 Trie 如下：

因为 Trie 是一个树形结构，其中肯定有表示节点的内部类Node。

我们发现，在上图所示的Trie中有些矩形是蓝色的，有些矩形是红色的，这表示什么意思呢？其实很简单，根据我们存储的单词可以发现，红色的矩形就表示单词结尾。那为什么要设置一个矩形是红色呢？

假设我们没有存储单词【dog】，那么这个时候第一个字母【g】下的矩形就应该是蓝色，但又因为存储了单词【doggy】，因此整个Trie储存的字母是不变的，仅仅是矩形颜色发生改变。在这个时候，如果我们查询单词【dog】是否存在于Trie中，就需要对Trie进行遍历。由于【g】下的矩形是蓝色，那么这个时候遍历出来的【dog】只是其他单词中的一部分，因此得到结果：单词【dog】是不存在于当前Trie中的。如果我们不设置蓝色或红色，那么这个Trie就无法正确判断是否存在单词【dog】。因此我们需要一个boolean类型的成员变量word表示当前节点是否为单词结尾。

---

在测试代码中，有这样一条语句trie.add("cat", 1);，其中的 cat 表示 key，1 表示value。那么对于Node的设计，成员变量 key 和 value 应该需要存在吧？😰

value 可以存在用来表示 1 ，但是需要注意的是，1 是储存在word == true的节点中的。但是 key 的存在就显得多余了。因为我们说过，我们在查询某个字符串时，经过一番节点查找其实就相当于获取了这个单词。比如我们查询单词【dog】，先找到 d ，又找到 0 ，最后找到 g ，同时遇到word == true的节点，表示前面扫描的字母组成了一个完整的单词，这个单词就是【dog】，既然如此，为了节约内存， key 是多余的。

那不要 key ，总得来个属性存储字母吧？

其实也不需要。当我们查询某个字符串时，是从根节点开始查找的，根节点下有很多个“叉”，同时根节点还是根据一个字符才找到一个特定的子节点。那应该怎么设计？

我们先假设字符只有 26 个英文字母，这样的话可以在Node类中设计一个数组，取名为children，数组长度为 26 ，里面存储着 26 个英文字母。这样的话，我们可以根据Node的children属性值来判断是否存在某个字母，然后得以判断是否存在这个字母对应的特定子节点。但很显然，这种做法的局限性非常大，只能存储英文字母，如果存储的是中文该咋办呢？难不成把所有中文存到一个数组里？😱

这个时候我们可以使用HashMap，HashMap的key就表示字符，value表示key对应的当前节点的子节点，这样就可以兼备所有字符情况了，同时，这也表明了我们的节点里是没有直接存储字符的，字符是存储在节点中HashMap的key中的。正因如此，在我们所画的Trie示意图中，字符是写在路径上的，并没有直接写在节点里。

通过上面一通分析，最后Node节点类设计如下：

private static class Node<V> {
    HashMap<Character, Node<V>> children;
    V value;
    boolean word; // 是否为单词结尾
}

2.3 清空、取值与包含

在Trie类中，暂时需要两个成员变量：

1、表示Tire的字符串数量 size

2、表示Trie的根节点 root

有了这两个成员变量后，我们就可以对一些方法进行实现了：

public class Trie<V> {
    private int size; // 元素（字符串）数量
    private Node<V> root;

    public int size() {
        return size;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public void clear() {
        size = 0;
        root = null;
    }
    
    // 省略其他代码
    ......
}

在上述代码中，我们完成了清空Trie的方法clear()，接下来准备编写另外两个方法。

V get(String key)方法表示在Trie中获取键为key的值（value）。

boolean contains(String key)方法表示判断当前Trie中是否存在键为key的数据。

我们发现这两个方法有一个共同点：需要根据key来获取节点，前者是获取节点的value，后者是判断节点是否存在。

既然如此，我们可以编写一个方法Node<V> node(String key)，这个方法可以根据传入的key在Trie中查找对应的节点。

为了方便获取某一节点的子节点，我们可以在内部类Node中编写一个方法，这个方法可以用来获取某个节点的子节点们。这样编写不仅可以使代码变得整洁，还可以减少node()方法中的判断。

private static class Node<V> {
    HashMap<Character, Node<V>> children;
    V value;
    boolean word; // 是否为单词结尾

    public HashMap<Character, Node<V>> getChildren() {
        return children == null ? (children = new HashMap<>()) : children;
    }
}

在node()方法中，还需要判断传入的参数是否为空，为此我们可以编写一个方法用来统一处理参数为空的情况：

private void keyCheck(String key) {
    if (key == null || key.length() == 0) {
        throw new IllegalArgumentException("key must not be empty");
    }
}

因为node()方法需要从根节点遍历Trie，因此我们在第一步要判断根节点是否为空。

node()方法代码如下：

private Node<V> node(String key) {
    if (root == null) return null;
    keyCheck(key);

    Node<V> node = root;
    int len = key.length();
    for (int i = 0; i < len; i++) {
        char c = key.charAt(i); // d o g
        node = node.getChildren().get(c);
        if (node == null) return null;
    }
    // 还得判断node是否是单词结尾
    return node.word ? node : null;
}

获取方法get()和包含方法contains()代码如下：

public V get(String key) {
    Node<V> node = node(key);
    return node == null ? null : node.value;
}

public boolean contains(String key) {
    return node(key) != null;
}

---

关于根节点的设置，我们还有另外一种方式，就是给成员变量root初始化：

private Node<V> root = new Node<>();

既然都进行了初始化，因此需要修改清空方法clear()，只清空Map，不设置为空，方便后续添加子节点：

public void clear() {
    size = 0;
    root.getChildren().clear();
}

初始化后还有一个好处，根节点永远不会为空，我们就可以删去node()方法中的非空判断：

private Node<V> node(String key) {
    //        if (root == null) return null;
    keyCheck(key);

    Node<V> node = root;
    int len = key.length();
    for (int i = 0; i < len; i++) {
        char c = key.charAt(i); // d o g
        node = node.getChildren().get(c);
        if (node == null) return null;
    }
    // 得判断node是否是单词结尾
    return node.word ? node : null;
}

2.4 添加字符串

添加元素的接口为：

public V add(String key, V value) {
    return null;
}

在实现这个接口时，我们首先要验证 key 是否为空。

添加元素的步骤如下：

1、获取根节点

2、获取 key 的长度

3、使用循环依次取出 key 中每个字符，然后从根节点向下遍历，判断其子节点是否是以 key 中的字符为HashMap的键

4、如果添加的元素已经存在，获取这个元素对应的 value，使用新 value 覆盖旧 value，最后返回旧的 value

5、如果添加的元素不存在，令元素最后一个字符对应的节点的 word 为 true，并使其 value 为传递进来的新 value，元素数量（字符串数量）加一，最后返回空

---

具体实现代码如下：

public V add(String key, V value) {
    keyCheck(key);

    Node<V> node = root;
    int len = key.length();
    for (int i = 0; i < len; i++) {
        char c = key.charAt(i); // d o g
        Node<V> childNode = node.getChildren().get(c);
        if (childNode == null) {
            childNode = new Node<>();
            node.getChildren().put(c, childNode);
        }
        node = childNode;
    }

    if (node.word) { // 已经存在这个单词
        V oldValue = node.value;
        node.value = value;
        return oldValue;
    }

    node.word = true;
    node.value = value;
    size++;
    return null;
}

2.5 前缀判断

所谓前缀判断就是判断当前Trie是否存在某字符串的前缀，注意区分与包含方法contains()的区别，包含方法是判断当前Trie中是否存在某一完整字符串。

举个例子，假设当前Trie存在一个字符串Student，我们可以使用前缀判断方法startsWith()判断当前Trie中是否存在前缀Stud，最终的结果返回true，我们也可以使用包含方法contains()判断当前Trie中是否存在完整字符串Student，最终返回结果是true。如果使用前缀判断当前Trie中是否存在前缀Student，结果也返回true，但是使用包含方法contains()判断当前Trie中是否存在字符串Stud，得到的返回结果是false。

前缀判断其实很简单，与node()方法极其相似，具体步骤如下：

1、对传入的前缀进行非空判断

2、获取根节点和前缀的长度

3、使用循环依次取出前缀中每个字符，然后从根节点向下遍历，判断其子节点是否存在，如果不存在就表示当前Trie中不存在传入的前缀，反之则存在

---

具体实现代码如下：

public boolean startsWith(String prefix) {
    keyCheck(prefix);

    Node<V> node = root;
    int len = prefix.length();
    for (int i = 0; i < len; i++) {
        char c = prefix.charAt(i);
        node = node.getChildren().get(c);
        if (node == null) return false;
    }
    return true;
}

2.6 代码调整

为了方便后续编写和节约空间，我们可以对代码进行调整，取消成员变量root的初始化，等用到根节点时才去创建根节点。

又回到最初的起点~🎵

private Node<V> root;

同时，在内部类Node中也没有必要获取子节点children时就创建一个，children也可以为空：

private static class Node<V> {
    HashMap<Character, Node<V>> children;
    V value;
    boolean word; // 是否为单词结尾
}

清空与添加

然后修改clear()方法：

public void clear() {
    size = 0;
    root = null;
}

在这种情况下，根节点root不是一定存在的，可能存在非空的情况，因此我们需要修改add()方法，当根节点为空时，创建根节点。同时因为取消了getChildren()方法， 添加方法add()修改如下：

public V add(String key, V value) {
    keyCheck(key);

    // 创建根节点
    if (root == null) {
        root = new Node<>();
    }

    Node<V> node = root;
    int len = key.length();
    for (int i = 0; i < len; i++) {
        char c = key.charAt(i); // d o g
        boolean emptyChild = node.children == null;
        Node<V> childNode = emptyChild ? null : node.children.get(c);
        if (childNode == null) {
            childNode = new Node<>();
            node.children = emptyChild ? new HashMap<>() : node.children;
            node.children.put(c, childNode);
        }
        node = childNode;
    }

    if (node.word) { // 已经存在这个单词
        V oldValue = node.value;
        node.value = value;
        return oldValue;
    }

    node.word = true;
    node.value = value;
    size++;
    return null;
}

获取节点与前缀判断

因为根节点不在初始化，导致根节点可能存在为空的情况，因此需要对node()方法和startsWith()方法进行修改。

同时还取消了getChildren()方法，对获取子节点的方式也需要进行修改。

因为node()方法和startsWith()方法的逻辑机器相似，存在代码复用的情况，我们可以对它俩一起进行修改：

public boolean startsWith(String prefix) {
        return node(prefix) != null;
    }

    private Node<V> node(String key) {
        keyCheck(key);

        Node<V> node = root;
        int len = key.length();
        for (int i = 0; i < len; i++) {
            // 注意这里的逻辑判断
            if (node == null || node.children == null || node.children.isEmpty()) return null;
            char c = key.charAt(i); // d o g
            node = node.children.get(c);
        }
     
//        return node.word ? node : null; // 得判断node是否是单词结尾
        return node; // 返回值可能为null
    }

取值与包含

因为node()方法的修改，get()与contains()也需要相应作出修改。

修改的主要原因是node()方法的返回值可能存在为null的情况，修改如下：

public V get(String key) {
    Node<V> node = node(key);
    return node != null && node.word ? node.value : null;
}

public boolean contains(String key) {
    Node<V> node = node(key);
    return node != null && node.word;
}

2.7 字符串删除

我们先理一下删除的逻辑：

1、获取传递的 key 对应的字符串的最后一个节点，我们将这个节点记为node

2、如果node为null或者不是单词结尾，不做任何处理，直接返回null

3、不满足第二步的情况下，node一定不为null且一定是单词结尾，这个时候必须要进行删除，我们先将size减一

4、如果node还有子节点，直接令其word为false，value为null

5、如果node没有子节点，就从node往回删

• 往回删除的时候遇到要被删除的节点有子节点或者其word为true时，就停止删除

6、删除完成后，记得返回node的value

---

为了能更好的编写删除方法reomve()，我们需要给内部类Node增加两个成员变量，并增加一个构造方法：

private static class Node<V> {
    Node<V> parent; // 父节点
    HashMap<Character, Node<V>> children;
    Character character; // 存入的字符
    V value;
    boolean word; // 是否为单词结尾

    public Node(Node<V> parent) {
        this.parent = parent;
    }
}

因为这个修改，我们需要对先前编写的add()方法进行修改：

public V add(String key, V value) {
    keyCheck(key);

    // 创建根节点
    if (root == null) {
        root = new Node<>(null);
    }

    Node<V> node = root;
    int len = key.length();
    for (int i = 0; i < len; i++) {
        char c = key.charAt(i); // d o g
        boolean emptyChild = node.children == null;
        Node<V> childNode = emptyChild ? null : node.children.get(c);
        if (childNode == null) {
            childNode = new Node<>(node);
            // 将字符存入Node成员变量中
            childNode.character = c;
            node.children = emptyChild ? new HashMap<>() : node.children;
            node.children.put(c, childNode);
        }
        node = childNode;
    }

    if (node.word) { // 已经存在这个单词
        V oldValue = node.value;
        node.value = value;
        return oldValue;
    }

    node.word = true;
    node.value = value;
    size++;
    return null;
}

最后，完成删除方法remove()的编写：

public V remove(String key) {
    // 找到最后一个节点
    Node<V> node = node(key);
    // 如果最后一个节点是空或不是单词结尾，不做任何处理
    if (node == null || !node.word) return null;
    size--;
    V oldValue = node.value;
    // 如果还有子节点
    if (node.children != null && !node.children.isEmpty()) {
        node.word = false;
        node.value = null;
        return oldValue;
    }
    // 如果没有子节点
    Node<V> parent = null;
    while ((parent = node.parent) != null) {
        parent.children.remove(node.character);
        if (parent.word || !parent.children.isEmpty()) break;
        node = parent;
    }
    return oldValue;
}

测试

到此，我们的Trie类就编写完了，可以进行测试，测试代码在文章最前面已经给出。

最终测试结果为：

2.8 总结

Tire的优点： 搜索前缀的效率主要跟前缀的长度有关 。

Trie的缺点：需要消耗大量的内存，因此有待改进。

更多Trie相关的数据结构和算法：

• Double-array Trie、Suffix Tree、Patricia Tree、Crit-bit Tree、AC 自动机