每日算法 — 使用java实现Trie字典树:前缀匹配与自动补全

算法背景

Trie(发音为 “try”),又称前缀树或字典树,是一种用于高效存储和检索字符串集合的树形数据结构。与哈希表相比,Trie 不仅能以 O(m) 时间完成单串查找,还能高效支持前缀查询、字典序遍历和自动补全等功能,是搜索引擎、拼写检查和路由表查找的核心数据结构。

核心思想

Trie 的每个节点代表一个字符,从根节点到任意节点的路径构成一个前缀。节点中维护 isEndOfWord 标记表示该节点是否为某个单词的结尾,同时通过子节点映射(数组或哈希表)实现字符到子树的跳转。

完整Java实现

import java.util.*;

public class TrieTree {

    /**
     * Trie节点定义
     */
    private static class TrieNode {
        // 使用HashMap支持任意字符,若仅 lowercase 可用大小为26的数组优化
        Map<Character, TrieNode> children;
        boolean isEndOfWord;      // 标记是否为单词结尾
        int prefixCount;          // 以该节点为前缀的单词数量(用于统计)

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

    private final TrieNode root;

    public TrieTree() {
        root = new TrieNode();
    }

    /**
     * 插入一个单词,时间复杂度 O(m),m为单词长度
     */
    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            node.children.putIfAbsent(ch, new TrieNode());
            node = node.children.get(ch);
            node.prefixCount++; // 统计经过该前缀的单词数
        }
        node.isEndOfWord = true;
    }

    /**
     * 查找单词是否存在,时间复杂度 O(m)
     */
    public boolean search(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            node = node.children.get(ch);
            if (node == null) return false;
        }
        return node.isEndOfWord;
    }

    /**
     * 判断是否有单词以给定前缀开头,时间复杂度 O(m)
     */
    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            node = node.children.get(ch);
            if (node == null) return false;
        }
        return true;
    }

    /**
     * 统计以prefix为前缀的单词数量
     */
    public int countPrefix(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            node = node.children.get(ch);
            if (node == null) return 0;
        }
        return node.prefixCount;
    }

    /**
     * 自动补全:返回以prefix为前缀的所有单词
     */
    public List<String> autoComplete(String prefix) {
        List<String> results = new ArrayList<>();
        TrieNode node = root;

        // 先定位到prefix的结尾节点
        for (char ch : prefix.toCharArray()) {
            node = node.children.get(ch);
            if (node == null) return results;
        }

        // DFS收集所有以该前缀开头的单词
        collectWords(node, new StringBuilder(prefix), results);
        return 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);
        }
    }

    /**
     * 删除单词(简易版,不清理无用节点)
     */
    public void delete(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            node = node.children.get(ch);
            if (node == null) return;
            node.prefixCount--;
        }
        node.isEndOfWord = false;
    }

    public static void main(String[] args) {
        TrieTree trie = new TrieTree();
        String[] words = {"java", "javascript", "javac", "python", "pytorch", "rust", "ruby"};

        System.out.println("插入单词: " + Arrays.toString(words));
        for (String w : words) trie.insert(w);

        System.out.println("search(\"java\"): " + trie.search("java"));
        System.out.println("search(\"jav\"): " + trie.search("jav"));
        System.out.println("startsWith(\"jav\"): " + trie.startsWith("jav"));
        System.out.println("startsWith(\"go\"): " + trie.startsWith("go"));

        System.out.println("\\n自动补全 'jav': " + trie.autoComplete("jav"));
        System.out.println("自动补全 'py': " + trie.autoComplete("py"));
        System.out.println("自动补全 'r': " + trie.autoComplete("r"));

        System.out.println("\\n以 'java' 为前缀的数量: " + trie.countPrefix("java"));
    }
}

复杂度分析

  • 插入:O(m) 时间,m 为单词长度。
  • 查找:O(m) 时间。
  • 前缀查询:O(m) 时间定位前缀节点,O(k) 时间收集 k 个结果。
  • 空间:最坏 O(字符集大小 × 总单词长度)。

应用场景

Trie 广泛应用于搜索引擎关键词提示、输入法自动补全、IP 路由最长前缀匹配、拼写检查词典、基因组序列分析等场景。通过将子节点数组替换为哈希表,可灵活支持 Unicode 字符集;在仅含小写字母的场景下,使用固定大小26的数组可获得更高的缓存命中率。