实战：Trie字典树（前缀匹配+自动补全）

说实话，字典树这东西，我第一次接触是在做一个搜索框自动补全功能的时候。当时用户输入一半，下面就要弹出候选词，用哈希表硬搞？效率太低了。后来翻到Trie这个数据结构，一拍大腿——这不就是为前缀匹配量身定做的吗？

今天咱们就手撸一个Trie树，把前缀匹配和自动补全这两个核心场景讲透。

什么是Trie树？

Trie树，也叫前缀树。说白了，就是一棵多叉树，每个节点代表一个字符。从根节点走到某个节点，路径上的字符连起来就是一个单词。

举个例子：插入"cat"、"car"、"dog"三个词，树长这样：

根 c d a o t r g cat car dog

每个节点存一个标志位，表示「这个节点是不是某个单词的结尾」。比如上图中带虚线圈的节点，就代表一个完整单词的结束。

核心数据结构

我习惯用数组或哈希表来存子节点。数组更快，但只适合小写字母这种有限字符集。哈希表更灵活，支持任意字符。

#include <string>
#include <unordered_map>
#include <vector>

class TrieNode {
public:
    std::unordered_map<char, TrieNode*> children;
    bool isEnd;  // 标记是否为单词结尾
    int count;   // 经过该节点的单词数量（可选）

    TrieNode() : isEnd(false), count(0) {}
};

class Trie {
private:
    TrieNode* root;

public:
    Trie() {
        root = new TrieNode();
    }

    // 插入单词
    void insert(const std::string& word) {
        TrieNode* node = root;
        for (char ch : word) {
            if (node->children.find(ch) == node->children.end()) {
                node->children[ch] = new TrieNode();
            }
            node = node->children[ch];
            node->count++;
        }
        node->isEnd = true;
    }

    // 搜索完整单词
    bool search(const std::string& word) {
        TrieNode* node = root;
        for (char ch : word) {
            if (node->children.find(ch) == node->children.end()) {
                return false;
            }
            node = node->children[ch];
        }
        return node->isEnd;
    }

    // 前缀匹配
    bool startsWith(const std::string& prefix) {
        TrieNode* node = root;
        for (char ch : prefix) {
            if (node->children.find(ch) == node->children.end()) {
                return false;
            }
            node = node->children[ch];
        }
        return true;
    }
};

关键点：插入时每经过一个节点，count++。这样后面做自动补全时，可以按热度排序。count越大，说明这个词被插入的次数越多，越应该排在前面。

自动补全：核心逻辑

自动补全分两步：

  1. 找到前缀对应的节点
  2. 从该节点开始DFS，收集所有以它为前缀的单词

代码实现：

// 自动补全：返回所有以prefix开头的单词
std::vector<std::string> autoComplete(const std::string& prefix) {
    std::vector<std::string> results;
    TrieNode* node = root;

    // 第一步：走到前缀末尾
    for (char ch : prefix) {
        if (node->children.find(ch) == node->children.end()) {
            return results;  // 前缀不存在，返回空
        }
        node = node->children[ch];
    }

    // 第二步：DFS收集所有单词
    std::string current = prefix;
    dfs(node, current, results);
    return results;
}

private:
void dfs(TrieNode* node, std::string& current, std::vector<std::string>& results) {
    if (node->isEnd) {
        results.push_back(current);
    }
    for (auto& [ch, child] : node->children) {
        current.push_back(ch);
        dfs(child, current, results);
        current.pop_back();  // 回溯
    }
}

我的小技巧：如果数据量很大，比如几十万词，DFS递归太深会爆栈。我一般改成迭代版，用栈模拟递归。或者限制最大返回数量，比如只返回前10个。

实战：带权重的自动补全

光有单词还不够，用户更希望看到「热门」的排在前面。我在项目中加了一个权重字段：

struct TrieNode {
    std::unordered_map<char, TrieNode*> children;
    bool isEnd;
    int weight;  // 权重，越大越靠前
    std::string word;  // 如果是结尾节点，存完整单词

    TrieNode() : isEnd(false), weight(0) {}
};

// 插入时更新权重
void insert(const std::string& word, int weight = 1) {
    TrieNode* node = root;
    for (char ch : word) {
        if (node->children.find(ch) == node->children.end()) {
            node->children[ch] = new TrieNode();
        }
        node = node->children[ch];
    }
    node->isEnd = true;
    node->weight += weight;  // 多次插入累加权重
    node->word = word;
}

// 带权重的自动补全
std::vector<std::pair<std::string, int>> autoCompleteWithWeight(
    const std::string& prefix, int topK = 10) {
    
    std::vector<std::pair<std::string, int>> results;
    TrieNode* node = root;

    for (char ch : prefix) {
        if (node->children.find(ch) == node->children.end()) {
            return results;
        }
        node = node->children[ch];
    }

    // 用优先队列（最小堆）维护topK
    std::priority_queue<std::pair<int, std::string>> pq;
    collectWords(node, prefix, pq);

    while (!pq.empty() && results.size() < topK) {
        results.push_back({pq.top().second, pq.top().first});
        pq.pop();
    }
    return results;
}

我曾经踩过的坑：权重累加时，如果同一个词被插入多次，权重会一直涨。但用户搜索「苹果」和「苹果手机」是两回事。后来我改成：每个用户会话内，同一个词只计一次权重。这样既反映了热度，又不会刷榜。

性能对比：Trie vs 其他方案

方案 插入复杂度 查询复杂度 内存占用 适用场景
Trie树 O(L) O(L) 较高（每个节点存指针） 前缀匹配、自动补全
哈希表 O(1) 平均 O(1) 平均 较低 精确查找
二分查找（排序数组） O(N) O(logN) 最低 静态数据
布隆过滤器 O(L) O(L) 极低 存在性判断（有误判）

你看，Trie在「前缀匹配」这个场景下，时间复杂度是O(L)，L是单词长度。哈希表虽然快，但没法做前缀匹配。排序数组可以用二分查找做前缀匹配，但插入成本太高。

优化思路

实际项目中，我一般会做几个优化：

  • 压缩节点：如果某个节点只有一个子节点，可以合并成一条边。这叫「压缩Trie」，内存能省不少。
  • 限制深度：自动补全时，DFS深度不要超过20层。用户输入的前缀越长，候选词越少，没必要搜太深。
  • 缓存热点：把前100个热门前缀的补全结果缓存起来，QPS能翻好几倍。

我个人的习惯：在Trie节点里加一个TrieNode* fail指针，实现AC自动机。这样不仅支持前缀匹配，还能做多模式串匹配——比如敏感词过滤。一个数据结构，两个用途，很划算。

嗯，Trie树的核心就这些。说白了，它就是空间换时间的典型代表。你想想看，如果不用Trie，每次前缀匹配都要遍历所有单词，那效率得多低？

代码写完了，建议你动手跑一跑。插入几千个词，试试自动补全的效果。你会发现，原来搜索引擎的「搜索建议」功能，底层就是这么个东西。

公众号:蓝海资料掘金营，微信deep3321

← 上一章 📖 返回目录 下一章 →