每日算法 — 使用java实现字典树:前缀树的高效字符串匹配与自动补全

数据结构引入

在处理大量字符串数据时,一个常见需求是快速判断某个字符串是否存在,或者查找具有相同前缀的所有字符串。传统哈希表虽然能在 $O(1)$ 时间内完成精确查找,却无法高效处理前缀匹配问题。字典树(Trie,又称前缀树或单词查找树)是一种专门用于字符串存储和检索的树形数据结构,它能在 $O(m)$ 时间内完成字符串查找($m$ 为字符串长度),并天然支持前缀匹配、自动补全等功能。

核心思想

字典树的核心思想是将字符串的每个字符作为树的一条边,从根节点到某一节点的路径就代表一个字符串前缀:

  • 根节点不包含字符,代表空字符串
  • 每个节点包含若干子节点,每条边对应一个字符
  • 从根到任意节点的路径上的字符拼接起来,就是该节点代表的字符串前缀
  • 在单词结束节点做特殊标记,表示从根到该节点是一条完整字符串

这种结构使得具有相同前缀的字符串共享路径,极大节省了存储空间,同时前缀查找只需沿着对应字符路径遍历即可。

完整Java实现

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 字典树(Trie / Prefix Tree)实现
 * 支持字符串插入、查找、前缀匹配、自动补全
 * 时间复杂度:插入和查找均为O(m),m为字符串长度
 */
public class Trie {

    /**
     * Trie节点定义
     * 使用HashMap存储子节点,支持任意字符集
     */
    private static class TrieNode {
        /** 子节点映射:字符 -> 子节点 */
        Map<Character, TrieNode> children;
        /** 标记该节点是否是一个完整单词的结尾 */
        boolean isEndOfWord;
        /** 记录以该节点结尾的单词数量(用于统计) */
        int wordCount;
        /** 经过该节点的单词前缀数量(用于统计) */
        int prefixCount;

        TrieNode() {
            this.children = new HashMap<>();
            this.isEndOfWord = false;
            this.wordCount = 0;
            this.prefixCount = 0;
        }
    }

    private final TrieNode root;

    public Trie() {
        this.root = new TrieNode();
    }

    /**
     * 向字典树中插入一个单词
     *
     * @param word 待插入的单词
     */
    public void insert(String word) {
        if (word == null || word.isEmpty()) {
            return;
        }

        TrieNode current = root;
        for (char ch : word.toCharArray()) {
            // 如果当前字符路径不存在,创建新节点
            current = current.children.computeIfAbsent(ch, k -> new TrieNode());
            // 增加前缀计数
            current.prefixCount++;
        }
        // 标记单词结尾
        current.isEndOfWord = true;
        current.wordCount++;
    }

    /**
     * 查找字典树中是否包含指定单词
     *
     * @param word 待查找的单词
     * @return 是否包含该完整单词
     */
    public boolean search(String word) {
        TrieNode node = searchNode(word);
        return node != null && node.isEndOfWord;
    }

    /**
     * 判断是否有单词以指定前缀开头
     *
     * @param prefix 前缀字符串
     * @return 是否存在以该前缀开头的单词
     */
    public boolean startsWith(String prefix) {
        return searchNode(prefix) != null;
    }

    /**
     * 查找指定字符串对应的终止节点
     *
     * @param str 目标字符串
     * @return 终止节点,不存在则返回null
     */
    private TrieNode searchNode(String str) {
        if (str == null || str.isEmpty()) {
            return null;
        }

        TrieNode current = root;
        for (char ch : str.toCharArray()) {
            TrieNode node = current.children.get(ch);
            if (node == null) {
                return null;
            }
            current = node;
        }
        return current;
    }

    /**
     * 统计字典树中以指定前缀开头的单词数量
     *
     * @param prefix 前缀
     * @return 以该前缀开头的单词数量
     */
    public int countWordsWithPrefix(String prefix) {
        TrieNode node = searchNode(prefix);
        return (node != null) ? countWords(node) : 0;
    }

    /**
     * 递归统计以某节点为根的子树中的单词数量
     */
    private int countWords(TrieNode node) {
        int count = node.isEndOfWord ? node.wordCount : 0;
        for (TrieNode child : node.children.values()) {
            count += countWords(child);
        }
        return count;
    }

    /**
     * 自动补全:获取以指定前缀开头的所有单词
     * 常用于搜索引擎建议、IDE代码补全
     *
     * @param prefix 前缀
     * @return 所有匹配单词的列表
     */
    public List<String> autoComplete(String prefix) {
        List<String> results = new ArrayList<>();
        TrieNode prefixNode = searchNode(prefix);

        if (prefixNode == null) {
            return results;
        }

        // 从prefix节点开始DFS收集所有单词
        collectWords(prefixNode, new StringBuilder(prefix), results);
        return results;
    }

    /**
     * DFS收集以当前节点为根的所有单词
     *
     * @param node    当前节点
     * @param prefix  当前构建的前缀
     * @param results 结果列表
     */
    private void collectWords(TrieNode node, StringBuilder prefix, List<String> results) {
        if (node.isEndOfWord) {
            results.add(prefix.toString());
        }

        for (Map.Entry<Character, TrieNode> entry : node.children.entrySet()) {
            prefix.append(entry.getKey());
            collectWords(entry.getValue(), prefix, results);
            prefix.deleteCharAt(prefix.length() - 1);
        }
    }

    /**
     * 从字典树中删除一个单词
     * 递归实现,清理不再需要的节点
     *
     * @param word 待删除的单词
     * @return 是否删除成功
     */
    public boolean delete(String word) {
        return delete(root, word, 0);
    }

    private boolean delete(TrieNode current, String word, int index) {
        if (index == word.length()) {
            // 到达单词结尾
            if (!current.isEndOfWord) {
                return false; // 单词不存在
            }
            current.isEndOfWord = false;
            current.wordCount--;
            // 如果没有子节点且不是其他单词前缀,可以删除
            return current.children.isEmpty();
        }

        char ch = word.charAt(index);
        TrieNode node = current.children.get(ch);
        if (node == null) {
            return false; // 单词不存在
        }

        boolean shouldDeleteChild = delete(node, word, index + 1);
        node.prefixCount--;

        if (shouldDeleteChild) {
            current.children.remove(ch);
            // 如果当前节点也不再是单词结尾且无其他子节点,返回true表示可以删除
            return !current.isEndOfWord && current.children.isEmpty();
        }

        return false;
    }

    /**
     * 查找字典树中最长公共前缀
     * 常用于URL路由匹配、DNA序列比对
     *
     * @return 所有单词的最长公共前缀
     */
    public String longestCommonPrefix() {
        StringBuilder prefix = new StringBuilder();
        TrieNode current = root;

        // 只要当前节点只有一个子节点且不是单词结尾,就继续深入
        while (current.children.size() == 1 && !current.isEndOfWord) {
            Map.Entry<Character, TrieNode> entry =
                current.children.entrySet().iterator().next();
            prefix.append(entry.getKey());
            current = entry.getValue();
        }

        return prefix.toString();
    }

    /**
     * 主方法:演示字典树的各项功能
     */
    public static void main(String[] args) {
        Trie trie = new Trie();

        // 批量插入单词
        String[] words = {
            "apple", "app", "application", "apply",
            "banana", "band", "bandana",
            "cat", "caterpillar", "category"
        };

        System.out.println("=== 插入单词 ===");
        for (String word : words) {
            trie.insert(word);
            System.out.println("插入: " + word);
        }

        // 查找测试
        System.out.println("\n=== 查找测试 ===");
        System.out.println("search('apple'): " + trie.search("apple"));
        System.out.println("search('app'): " + trie.search("app"));
        System.out.println("search('appl'): " + trie.search("appl"));
        System.out.println("startsWith('app'): " + trie.startsWith("app"));
        System.out.println("startsWith('ban'): " + trie.startsWith("ban"));
        System.out.println("startsWith('xyz'): " + trie.startsWith("xyz"));

        // 自动补全测试
        System.out.println("\n=== 自动补全 ===");
        System.out.println("autoComplete('app'): " + trie.autoComplete("app"));
        System.out.println("autoComplete('ban'): " + trie.autoComplete("ban"));
        System.out.println("autoComplete('cat'): " + trie.autoComplete("cat"));

        // 最长公共前缀
        System.out.println("\n=== 最长公共前缀 ===");
        System.out.println("Longest Common Prefix: '" + trie.longestCommonPrefix() + "'");

        // 删除测试
        System.out.println("\n=== 删除测试 ===");
        System.out.println("删除 'app': " + trie.delete("app"));
        System.out.println("删除后 search('app'): " + trie.search("app"));
        System.out.println("删除后 search('apple'): " + trie.search("apple"));
        System.out.println("autoComplete('app') after delete: " + trie.autoComplete("app"));
    }
}

复杂度分析

操作 时间复杂度 空间复杂度 说明
插入 $O(m)$ $O(m)$ $m$ 为字符串长度,最坏情况需创建 $m$ 个新节点
查找 $O(m)$ $O(1)$ 只需遍历 $m$ 个字符对应的边
前缀匹配 $O(m)$ $O(k)$ $k$ 为匹配结果数量
删除 $O(m)$ $O(1)$ 递归清理无用节点

应用场景

字典树在以下场景发挥关键作用:

  • 搜索引擎自动补全:用户输入前缀时实时返回匹配建议,Trie的前缀特性使其天然适合
  • 拼写检查与纠错:快速判断单词是否存在,结合编辑距离可扩展为纠错功能
  • IP路由最长前缀匹配:路由器根据目的IP查找最长匹配前缀决定转发路径
  • DNA序列分析:基因组数据本质上是{A,C,G,T}上的字符串,Trie可高效存储和比对
  • Word Filter敏感词过滤:构建敏感词Trie,对输入文本逐字符匹配,一旦发现完整敏感词即拦截

字典树通过字符级别的路径共享,将字符串集合转化为树形结构,使前缀相关操作达到理论最优的线性时间复杂度。作为字符串处理领域的基石数据结构,它在编译器词法分析、数据库索引、生物信息学等领域都有不可替代的地位。