每日算法 — 使用java实现Trie前缀树:高效字符串存储与自动补全

Trie(发音同”try”),又称前缀树或字典树,是一种树形数据结构,专门用于高效存储和检索字符串集合。与普通二叉搜索树不同,Trie的每个节点代表一个字符串前缀,从根到任意节点的路径即为一个前缀。Trie在搜索引擎自动补全、拼写检查、IP路由查找和前缀匹配等场景中有着广泛应用。本文将用Java实现一个功能完整的Trie,支持插入、搜索、前缀匹配与自动补全。

Trie核心结构

Trie的每个节点包含:
– 子节点数组/映射:通常使用数组(固定字符集)或哈希映射实现
– isEnd标记:表示从根到该节点是否构成一个完整单词
– 计数器(可选):记录经过该节点的单词数量,支持前缀频次统计

关键性质:
– 查找和插入时间复杂度均为O(L),L为字符串长度
– 与集合中单词数量无关,只与单词长度有关
– 空间复杂度较高,存在大量公共前缀时可显著节省空间

Java完整实现

import java.util.*;

/**
 * Trie前缀树实现
 * 支持:插入、精确搜索、前缀搜索、自动补全、删除
 */
public class Trie {

    // Trie节点
    static class TrieNode {
        // 使用HashMap存储子节点,支持任意字符(包括中文)
        Map<Character, TrieNode> children;
        // 标记是否为单词结尾
        boolean isEndOfWord;
        // 经过该节点的单词数量(用于前缀频次统计)
        int passCount;
        // 以该节点结尾的单词数量(支持重复单词)
        int endCount;

        TrieNode() {
            children = new HashMap<>();
            isEndOfWord = false;
            passCount = 0;
            endCount = 0;
        }
    }

    private TrieNode root;
    private int size; // 存储的不同单词数量

    public Trie() {
        root = new TrieNode();
        size = 0;
    }

    /**
     * 插入一个单词
     * 时间复杂度: O(L),L为单词长度
     */
    public void insert(String word) {
        if (word == null || word.isEmpty()) return;

        TrieNode curr = root;
        for (char ch : word.toCharArray()) {
            curr.passCount++;
            curr = curr.children.computeIfAbsent(ch, k -> new TrieNode());
        }
        curr.passCount++;
        curr.endCount++;
        if (!curr.isEndOfWord) {
            curr.isEndOfWord = true;
            size++;
        }
    }

    /**
     * 精确搜索:判断单词是否在Trie中
     */
    public boolean search(String word) {
        TrieNode node = findNode(word);
        return node != null && node.isEndOfWord;
    }

    /**
     * 前缀搜索:判断是否有单词以给定前缀开头
     */
    public boolean startsWith(String prefix) {
        return findNode(prefix) != null;
    }

    /**
     * 获取以prefix为前缀的单词数量
     */
    public int countPrefix(String prefix) {
        TrieNode node = findNode(prefix);
        return node == null ? 0 : node.passCount;
    }

    /**
     * 辅助方法:找到代表给定字符串的最后一个节点
     */
    private TrieNode findNode(String str) {
        if (str == null || str.isEmpty()) return null;

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

    /**
     * 自动补全:获取以prefix开头的所有单词
     * @param prefix 前缀
     * @param maxResults 最多返回的结果数
     */
    public List<String> autoComplete(String prefix, int maxResults) {
        List<String> results = new ArrayList<>();
        if (prefix == null) return results;

        // 先找到prefix对应的节点
        TrieNode prefixNode = root;
        for (char ch : prefix.toCharArray()) {
            prefixNode = prefixNode.children.get(ch);
            if (prefixNode == null) return results; // 无前缀匹配
        }

        // 从该节点开始DFS收集所有单词
        StringBuilder path = new StringBuilder(prefix);
        dfsCollect(prefixNode, path, results, maxResults);

        return results;
    }

    /**
     * DFS收集所有单词
     */
    private void dfsCollect(TrieNode node, StringBuilder path, 
                           List<String> results, int maxResults) {
        if (results.size() >= maxResults) return;

        if (node.isEndOfWord) {
            results.add(path.toString());
        }

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

    /**
     * 删除一个单词
     * 返回是否成功删除
     */
    public boolean delete(String word) {
        return deleteHelper(root, word, 0);
    }

    private boolean deleteHelper(TrieNode curr, String word, int index) {
        if (index == word.length()) {
            if (!curr.isEndOfWord) return false; // 单词不存在
            curr.endCount--;
            if (curr.endCount == 0) {
                curr.isEndOfWord = false;
                size--;
            }
            return curr.children.isEmpty(); // 是否可以删除该节点
        }

        char ch = word.charAt(index);
        TrieNode child = curr.children.get(ch);
        if (child == null) return false;

        boolean shouldDeleteChild = deleteHelper(child, word, index + 1);

        if (shouldDeleteChild) {
            curr.children.remove(ch);
            // 如果当前节点不再被需要(不是单词结尾且无其他子节点)
            return !curr.isEndOfWord && curr.children.isEmpty();
        }

        return false;
    }

    /**
     * 获取Trie中存储的不同单词数量
     */
    public int size() {
        return size;
    }

    /**
     * 检查Trie是否为空
     */
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 获取Trie中所有单词
     */
    public List<String> getAllWords() {
        List<String> words = new ArrayList<>();
        dfsCollect(root, new StringBuilder(), words, Integer.MAX_VALUE);
        return words;
    }

    // ==================== 应用示例 ====================

    /**
     * 搜索引擎自动补全示例
     */
    public static void searchEngineDemo() {
        Trie trie = new Trie();

        // 模拟搜索日志中的热门查询
        String[] queries = {
            "java", "javascript", "java tutorial", "java spring",
            "python", "python django", "python flask",
            "algorithm", "algorithm design", "algorithm analysis",
            "data structure", "database", "database design",
            "machine learning", "deep learning",
            "web development", "web design", "web security"
        };

        for (String q : queries) {
            trie.insert(q);
        }

        System.out.println("=== 搜索引擎自动补全演示 ===");

        String[] prefixes = {"ja", "py", "web", "alg", "data"};
        for (String prefix : prefixes) {
            System.out.println("\n输入: \"" + prefix + "\"");
            List<String> suggestions = trie.autoComplete(prefix, 5);
            System.out.println("建议: " + suggestions);
        }
    }

    /**
     * 拼写检查与纠错示例
     */
    public static void spellCheckDemo() {
        Trie dictionary = new Trie();

        // 加载词典
        String[] words = {
            "algorithm", "algorithms", "algoritm",  // 注意:故意包含拼写错误
            "structure", "structured", "structures",
            "programming", "program", "programs",
            "computer", "computing", "compute"
        };

        for (String w : words) {
            dictionary.insert(w);
        }

        System.out.println("\n=== 拼写检查演示 ===");
        String[] checkWords = {"algoritm", "strcture", "computr", "programing"};

        for (String word : checkWords) {
            boolean exists = dictionary.search(word);
            System.out.println("\"" + word + "\" 在词典中: " + exists);
            if (!exists) {
                // 简单纠错:查找编辑距离为1的候选
                List<String> candidates = findCandidates(dictionary, word);
                System.out.println("  建议修正: " + candidates);
            }
        }
    }

    /**
     * 简单候选生成:删除一个字符后检查是否在词典中
     */
    private static List<String> findCandidates(Trie trie, String word) {
        List<String> candidates = new ArrayList<>();
        for (int i = 0; i < word.length(); i++) {
            String candidate = word.substring(0, i) + word.substring(i + 1);
            if (trie.search(candidate)) {
                candidates.add(candidate);
            }
        }
        return candidates;
    }

    public static void main(String[] args) {
        searchEngineDemo();
        spellCheckDemo();

        // 基础功能测试
        System.out.println("\n=== 基础功能测试 ===");
        Trie trie = new Trie();
        trie.insert("apple");
        trie.insert("app");
        trie.insert("application");
        trie.insert("apply");

        System.out.println("搜索 'app': " + trie.search("app"));
        System.out.println("搜索 'apple': " + trie.search("apple"));
        System.out.println("前缀 'app': " + trie.startsWith("app"));
        System.out.println("自动补全 'app': " + trie.autoComplete("app", 10));
        System.out.println("单词总数: " + trie.size());
    }
}

复杂度分析

操作 时间复杂度 空间复杂度 说明
插入 O(L) O(L) L为单词长度
精确搜索 O(L) O(1) 仅需遍历
前缀搜索 O(L) O(1) 仅需遍历前缀
自动补全 O(L + N) O(N) N为匹配结果数
删除 O(L) O(1) 可能释放节点

Trie变体与优化

  1. 压缩Trie(Radix Tree):合并只有一个子节点的链,减少空间占用
  2. 后缀树:存储所有后缀,支持子串查询
  3. 双数组Trie:用两个数组表示Trie,空间效率极高
  4. AC自动机:在Trie上构建失败指针,支持多模式串匹配
  5. Xor Trie:用于最大异或子集等问题