每日算法 — 使用java实现字典树:前缀树高效检索与自动补全系统

引言

在信息检索、输入法联想、搜索引擎建议等场景中,快速查找具有相同前缀的字符串集合是一个核心需求。字典树(Trie,又称前缀树)正是为解决这类问题而生的经典数据结构。它以空间换时间,将字符串的公共前缀共享存储,使得前缀查询的时间复杂度仅与目标串长度相关。本文将使用Java从零实现一个完整的字典树,涵盖基础插入删除、前缀搜索、自动补全、模式匹配等高级功能。

数据结构原理

字典树是一棵多叉树,每条边代表一个字符。从根节点到任意节点的路径上的字符序列构成一个前缀。每个节点维护以下信息:

  • 子节点映射:字符到子节点的映射,通常用数组或哈希表实现
  • 是否为单词结尾:标记从根到该节点是否构成一个完整单词
  • 经过次数/单词计数:可选,用于统计前缀出现频率或自动补全排序

核心操作

  • 插入:从根出发,沿字符路径下行,不存在则创建新节点,最后标记结尾
  • 搜索:沿字符路径下行,中途断链则失败,到达后检查结尾标记
  • 前缀匹配:仅需沿路径下行,无需检查结尾标记
  • 删除:找到目标节点,清除结尾标记,若该节点无子节点则回溯清理空分支

Java实现

import java.util.*;

/**
 * 字典树节点
 */
class TrieNode {
    // 使用HashMap存储子节点,支持任意字符集
    Map<Character, TrieNode> children;
    // 标记是否为某个单词的结尾
    boolean isEndOfWord;
    // 以该节点为结尾的单词数量(支持重复插入)
    int wordCount;
    // 经过该节点的单词数量(用于前缀频率统计)
    int passCount;

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

/**
 * 字典树(前缀树)完整实现
 * 支持:插入、删除、精确查找、前缀查找、自动补全、通配符匹配
 */
public class Trie {
    private final TrieNode root;
    private int size; // 存储的不重复单词数量

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

    /**
     * 插入一个单词,时间复杂度O(L),L为单词长度
     */
    public void insert(String word) {
        if (word == null || word.isEmpty()) return;
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            node.children.putIfAbsent(ch, new TrieNode());
            node = node.children.get(ch);
            node.passCount++;
        }
        if (!node.isEndOfWord) {
            size++;
        }
        node.isEndOfWord = true;
        node.wordCount++;
    }

    /**
     * 精确查找单词是否存在
     */
    public boolean search(String word) {
        TrieNode node = findNode(word);
        return node != null && node.isEndOfWord;
    }

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

    /**
     * 获取指定前缀的经过次数(有多少单词包含此前缀)
     */
    public int prefixCount(String prefix) {
        TrieNode node = findNode(prefix);
        return node != null ? node.passCount : 0;
    }

    /**
     * 获取以指定前缀开头的所有单词(自动补全)
     */
    public List<String> autoComplete(String prefix) {
        List<String> results = new ArrayList<>();
        TrieNode node = findNode(prefix);
        if (node == null) return results;
        collectWords(node, new StringBuilder(prefix), results);
        return results;
    }

    /**
     * 获取以指定前缀开头的最多limit个单词,按字典序排列
     */
    public List<String> autoComplete(String prefix, int limit) {
        List<String> results = autoComplete(prefix);
        return results.size() > limit ? results.subList(0, limit) : results;
    }

    /**
     * 通配符搜索:'.'匹配任意单个字符
     */
    public boolean searchWithWildcard(String pattern) {
        return searchWildcard(root, pattern, 0);
    }

    /**
     * 删除一个单词,若该单词不存在则不做任何事
     */
    public void delete(String word) {
        if (word == null || word.isEmpty()) return;
        delete(root, word, 0);
    }

    /**
     * 获取字典树中存储的不重复单词总数
     */
    public int size() {
        return size;
    }

    /**
     * 获取所有存储的单词
     */
    public List<String> getAllWords() {
        List<String> words = new ArrayList<>();
        collectWords(root, new StringBuilder(), words);
        return words;
    }

    // ==================== 私有辅助方法 ====================

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

    private void collectWords(TrieNode node, StringBuilder prefix, List<String> results) {
        if (node.isEndOfWord) {
            results.add(prefix.toString());
        }
        // 按字符排序遍历,保证字典序
        List<Character> keys = new ArrayList<>(node.children.keySet());
        Collections.sort(keys);
        for (char ch : keys) {
            prefix.append(ch);
            collectWords(node.children.get(ch), prefix, results);
            prefix.deleteCharAt(prefix.length() - 1);
        }
    }

    private boolean searchWildcard(TrieNode node, String pattern, int index) {
        if (index == pattern.length()) {
            return node.isEndOfWord;
        }
        char ch = pattern.charAt(index);
        if (ch == '.') {
            for (TrieNode child : node.children.values()) {
                if (searchWildcard(child, pattern, index + 1)) {
                    return true;
                }
            }
            return false;
        } else {
            TrieNode child = node.children.get(ch);
            if (child == null) return false;
            return searchWildcard(child, pattern, index + 1);
        }
    }

    private boolean delete(TrieNode node, String word, int index) {
        if (index == word.length()) {
            if (!node.isEndOfWord) return false; // 单词不存在
            node.isEndOfWord = false;
            node.wordCount--;
            size--;
            return node.children.isEmpty(); // 是否可以删除该节点
        }
        char ch = word.charAt(index);
        TrieNode child = node.children.get(ch);
        if (child == null) return false;

        boolean shouldDeleteChild = delete(child, word, index + 1);
        if (shouldDeleteChild) {
            node.children.remove(ch);
            return !node.isEndOfWord && node.children.isEmpty();
        }
        child.passCount--; // 减少经过计数
        return false;
    }
}

/**
 * 基于数组优化的字典树(仅适用于小写英文字母)
 * 使用固定大小数组代替HashMap,访问更快,内存更紧凑
 */
class ArrayTrie {
    private static final int ALPHABET_SIZE = 26;

    static class Node {
        Node[] children = new Node[ALPHABET_SIZE];
        boolean isEndOfWord;
    }

    private final Node root;

    public ArrayTrie() {
        root = new Node();
    }

    public void insert(String word) {
        Node node = root;
        for (char ch : word.toCharArray()) {
            int idx = ch - 'a';
            if (node.children[idx] == null) {
                node.children[idx] = new Node();
            }
            node = node.children[idx];
        }
        node.isEndOfWord = true;
    }

    public boolean search(String word) {
        Node node = root;
        for (char ch : word.toCharArray()) {
            int idx = ch - 'a';
            if (node.children[idx] == null) return false;
            node = node.children[idx];
        }
        return node.isEndOfWord;
    }
}

/**
 * 字典树应用示例:拼写检查与自动补全系统
 */
public class TrieApplication {
    public static void main(String[] args) {
        Trie trie = new Trie();

        // 模拟词典加载
        String[] dictionary = {
            "java", "javascript", "javac", "javadoc",
            "algorithm", "algorithms", "align", "alignment",
            "tree", "trie", "try", "type",
            "search", "searches", "searching", "searched"
        };
        for (String word : dictionary) {
            trie.insert(word);
        }

        System.out.println("词典大小: " + trie.size());
        System.out.println("查找 'java': " + trie.search("java"));
        System.out.println("查找 'jav': " + trie.search("jav"));
        System.out.println("前缀 'jav' 是否存在: " + trie.startsWith("jav"));

        System.out.println("\n自动补全 'sear':");
        List<String> completions = trie.autoComplete("sear");
        completions.forEach(System.out::println);

        System.out.println("\n通配符搜索 'j.v.':");
        System.out.println("匹配 'java': " + trie.searchWithWildcard("j.va"));
        System.out.println("匹配 'javascript': " + trie.searchWithWildcard("j.vascript"));

        System.out.println("\n删除 'java' 后:");
        trie.delete("java");
        System.out.println("查找 'java': " + trie.search("java"));
        System.out.println("查找 'javascript': " + trie.search("javascript"));

        // 前缀频率统计
        System.out.println("\n前缀 'al' 的经过次数: " + trie.prefixCount("al"));
    }
}

进阶:压缩字典树(Radix Tree)

当存储的单词存在大量公共前缀时,标准字典树的节点数可能膨胀。压缩字典树将单分支链合并为一条边,边标签变为字符串片段,显著减少节点数量。

/**
 * 压缩字典树节点
 */
class RadixNode {
    Map<String, RadixNode> children = new HashMap<>();
    boolean isEndOfWord;
    // 可选:存储边标签的公共前缀长度等额外信息
}

复杂度分析

操作 时间复杂度 空间复杂度
插入 O(L) O(L) 最坏情况
精确查找 O(L) O(1)
前缀查找 O(L) O(1)
自动补全(输出K个结果) O(L + K×M) O(K×M)
通配符搜索 O(26^L) 最坏 O(L) 递归栈
删除 O(L) O(1)

其中L为查询串长度,M为结果串平均长度。

工程应用场景

  1. 搜索引擎建议框:用户输入前缀,毫秒级返回Top-K热门搜索词
  2. 输入法联想:基于拼音前缀匹配候选汉字词组
  3. IDE代码补全:前缀匹配变量名、类名、方法名
  4. 路由表查找:IP路由最长前缀匹配(可用压缩字典树优化)
  5. DNA序列分析:基因组数据中频繁出现的前缀统计

总结

字典树通过共享公共前缀,将字符串集合的组织从”线性列表”升级为”树形索引”,使前缀相关操作从O(N)降至O(L)。掌握字典树的核心在于理解”路径即前缀”的思想,以及节点状态(子节点引用、结尾标记)的正确维护。本文提供的Java实现涵盖基础与高级功能,可直接应用于拼写检查、自动补全等实际系统,也可作为学习更复杂字符串数据结构(如后缀树、AC自动机)的基石。