实现 Trie (前缀树) | 秋招知识库

题目链接 · 灵神原题解（署名来源）

思路

“若我俩都走过同样的路，则我俩心意相通。”

假设字符串里面只有 $a$ 和 $b$ 两种字符。

从左到右遍历字符串，把 $a$ 视作往左走，把 $b$ 视作往右走。怎么记录路径？怎么区分不同的路径？用二叉树记录。

$insert$ ：例如先插入字符串 $aab$ ，相当于生成了一条移动方向为「左-左-右」的路径。标记最后一个节点为终止节点。再插入字符串 $aabb$ ，相当于生成了一条移动方向为「左-左-右-右」的路径。标记最后一个节点为终止节点。
$search$ ：例如查找字符串 $aab$ ，相当于查找二叉树中是否存在一条移动方向为「左-左-右」的路径，且最后一个节点是终止节点。
$startsWith$ ：例如查找前缀 $aa$ ，相当于查找二叉树中是否存在一条移动方向为「左-左」的路径，无其他要求。

推广到 $26$ 种字母，其实就是一棵 $26$ 叉树，对于 $26$ 叉树的每个节点，可以用哈希表，或者长为 $26$ 的数组来存储子节点。

算法

初始化：创建一棵 $26$ 叉树，一开始只有一个根节点 $root$ 。 $26$ 叉树的每个节点包含一个长为 $26$ 的儿子节点列表 $son$ ，以及一个布尔值 $end$ ，表示是否为终止节点。
$insert$ ：
1. 遍历字符串 $word$ ，同时用一个变量 $cur$ 表示当前在 $26$ 叉树的哪个节点，初始值为 $root$ 。
2. 如果 $word [i]$ 不是 $cur$ 的儿子，那么创建一个新的节点 $node$ 作为 $cur$ 的儿子。如果 $word [i] = a$ ，那么把 $node$ 记录到 $cur$ 的 $son [0]$ 中。如果 $word [i] = b$ ，那么把 $node$ 记录到 $cur$ 的 $son [1]$ 中。依此类推。
3. 更新 $cur$ 为儿子列表中的相应节点。
4. 遍历结束，把 $cur$ 的 $end$ 标记为 $true$ 。
$search$ 和 $startsWith$ 可以复用同一个函数 $find$ ：
1. 遍历字符串 $word$ ，同时用一个变量 $cur$ 表示当前在 $26$ 叉树的哪个节点，初始值为 $root$ 。
2. 如果 $word [i]$ 不是 $cur$ 的儿子，返回 $0$ 。 $search$ 和 $startsWith$ 收到 $0$ 之后返回 $false$ 。
3. 更新 $cur$ 为儿子列表中的相应节点。
4. 遍历结束，如果 $cur$ 的 $end$ 是 $false$ ，返回 $1$ ，否则返回 $2$ 。 $search$ 如果收到的是 $2$ ，返回 $true$ ，否则返回 $false$ 。 $startsWith$ 如果收到的是非 $0$ 数字，返回 $true$ ，否则返回 $false$ 。

Python 解法1C++ 解法1

python

class Node:
    __slots__ = 'son', 'end'

    def __init__(self):
        self.son = {}
        self.end = False

class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, word: str) -> None:
        cur = self.root
        for c in word:
            if c not in cur.son:  # 无路可走？
                cur.son[c] = Node()  # 那就造路！
            cur = cur.son[c]
        cur.end = True

    def find(self, word: str) -> int:
        cur = self.root
        for c in word:
            if c not in cur.son:  # 道不同，不相为谋
                return 0
            cur = cur.son[c]
        # 走过同样的路（2=完全匹配，1=前缀匹配）
        return 2 if cur.end else 1

    def search(self, word: str) -> bool:
        return self.find(word) == 2

    def startsWith(self, prefix: str) -> bool:
        return self.find(prefix) != 0

cpp

// C++ 版待补充

Python 解法2C++ 解法2

python

class Node:
    __slots__ = 'son', 'end'

    def __init__(self):
        self.son = [None] * 26
        self.end = False

class Trie:
    def __init__(self):
        self.root = Node()

    def insert(self, word: str) -> None:
        cur = self.root
        for c in word:
            c = ord(c) - ord('a')
            if cur.son[c] is None:  # 无路可走？
                cur.son[c] = Node()  # 那就造路！
            cur = cur.son[c]
        cur.end = True

    def find(self, word: str) -> int:
        cur = self.root
        for c in word:
            c = ord(c) - ord('a')
            if cur.son[c] is None:  # 道不同，不相为谋
                return 0
            cur = cur.son[c]
        # 走过同样的路（2=完全匹配，1=前缀匹配）
        return 2 if cur.end else 1

    def search(self, word: str) -> bool:
        return self.find(word) == 2

    def startsWith(self, prefix: str) -> bool:
        return self.find(prefix) != 0

cpp

// C++ 版待补充

C++ 解法3

cpp

struct Node {
    Node* son[26]{};
    bool end = false;
};

class Trie {
    Node* root = new Node();

    int find(string word) {
        Node* cur = root;
        for (char c : word) {
            c -= 'a';
            if (cur->son[c] == nullptr) { // 道不同，不相为谋
                return 0;
            }
            cur = cur->son[c];
        }
        // 走过同样的路（2=完全匹配，1=前缀匹配）
        return cur->end ? 2 : 1;
    }

    void destroy(Node* node) {
        if (node == nullptr) {
            return;
        }
        for (Node* son : node->son) {
            destroy(son);
        }
        delete node;
    }

public:
    ~Trie() {
        destroy(root);
    }

    void insert(string word) {
        Node* cur = root;
        for (char c : word) {
            c -= 'a';
            if (cur->son[c] == nullptr) { // 无路可走？
                cur->son[c] = new Node(); // new 出来！
            }
            cur = cur->son[c];
        }
        cur->end = true;
    }

    bool search(string word) {
        return find(word) == 2;
    }

    bool startsWith(string prefix) {
        return find(prefix) != 0;
    }
};

复杂度分析

时间复杂度：初始化为 $O (1)$ ， $insert$ 为 $O (n | Σ |)$ ，其余为 $O (n)$ ，其中 $n$ 是 $word$ 的长度， $| Σ | = 26$ 是字符集合的大小。注意创建一个节点需要 $O (| Σ |)$ 的时间（如果用的是数组）。
空间复杂度： $O (q n | Σ |)$ 。其中 $q$ 是 $insert$ 的调用次数。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

思路 ​

算法 ​

复杂度分析 ​

分类题单 ​

思路

算法

复杂度分析

分类题单