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

引言

字典树(Trie,又称前缀树或单词查找树)是一种专门用于处理字符串集合的树形数据结构。其核心优势在于利用字符串的公共前缀来减少查询时间,使得插入和查找的时间复杂度仅与字符串长度相关,而与集合中元素总数无关。Trie树在搜索引擎自动补全、拼写检查、IP路由前缀匹配等场景中有着不可替代的作用。

节点结构设计

Trie的每个节点包含若干子节点指针(通常用数组或哈希表实现)以及一个标记位,表示从根到该节点的路径是否构成一个完整单词。

import java.util.*;

public class TrieNode {
    // 使用HashMap存储子节点,支持任意字符集
    Map<Character, TrieNode> children;
    boolean isEndOfWord; // 标记是否为一个完整单词的结尾
    int wordCount;       // 经过该节点的单词数量(用于前缀统计)

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

核心操作实现

public class Trie {
    private TrieNode root;

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

    // 插入一个单词
    public void insert(String word) {
        TrieNode current = root;
        for (char ch : word.toCharArray()) {
            current.children.putIfAbsent(ch, new TrieNode());
            current = current.children.get(ch);
            current.wordCount++; // 统计经过该节点的单词数
        }
        current.isEndOfWord = true;
    }

    // 精确查找:判断单词是否存在于Trie中
    public boolean search(String word) {
        TrieNode node = searchNode(word);
        return node != null && node.isEndOfWord;
    }

    // 前缀查找:判断是否存在以prefix开头的单词
    public boolean startsWith(String prefix) {
        return searchNode(prefix) != null;
    }

    // 统计以prefix为前缀的单词数量
    public int countWordsWithPrefix(String prefix) {
        TrieNode node = searchNode(prefix);
        return node == null ? 0 : node.wordCount;
    }

    // 辅助方法:找到对应字符串的最后一个节点
    private TrieNode searchNode(String str) {
        TrieNode current = root;
        for (char ch : str.toCharArray()) {
            TrieNode node = current.children.get(ch);
            if (node == null) return null;
            current = node;
        }
        return current;
    }
}

自动补全功能实现

利用Trie的前缀特性,可以实现高效的自动补全搜索:输入前缀后,DFS遍历以该前缀结尾节点为根的所有子树,收集完整单词。

    // 自动补全:返回所有以给定前缀开头的单词
    public List<String> autoComplete(String prefix) {
        List<String> results = new ArrayList<>();
        TrieNode prefixNode = searchNode(prefix);
        if (prefixNode == null) return results;

        StringBuilder currentWord = new StringBuilder(prefix);
        dfsCollectWords(prefixNode, currentWord, results);
        return results;
    }

    // 深度优先搜索收集所有单词
    private void dfsCollectWords(TrieNode node, StringBuilder path, List<String> results) {
        if (node.isEndOfWord) {
            results.add(path.toString());
        }
        for (Map.Entry<Character, TrieNode> entry : node.children.entrySet()) {
            path.append(entry.getKey());
            dfsCollectWords(entry.getValue(), path, results);
            path.deleteCharAt(path.length() - 1); // 回溯
        }
    }

带权Trie:实现搜索热词排序

在实际搜索引擎中,补全结果应按搜索频率排序。可扩展节点结构,在叶子节点存储权重,DFS时按权重排序输出。

public class WeightedTrieNode {
    Map<Character, WeightedTrieNode> children = new HashMap<>();
    boolean isEndOfWord;
    int frequency; // 单词搜索频率
}

public class WeightedTrie {
    private WeightedTrieNode root;

    public WeightedTrie() { root = new WeightedTrieNode(); }

    public void insert(String word, int freq) {
        WeightedTrieNode current = root;
        for (char ch : word.toCharArray()) {
            current.children.putIfAbsent(ch, new WeightedTrieNode());
            current = current.children.get(ch);
        }
        current.isEndOfWord = true;
        current.frequency = freq;
    }

    // 获取前缀补全结果,按频率降序排列
    public List<String> autoCompleteSorted(String prefix, int limit) {
        WeightedTrieNode prefixNode = findNode(prefix);
        if (prefixNode == null) return Collections.emptyList();

        PriorityQueue<WeightedTrieNode> pq = new PriorityQueue<>(
            (a, b) -> b.frequency - a.frequency
        );
        collectAllWords(prefixNode, prefix, pq);

        List<String> result = new ArrayList<>();
        while (!pq.isEmpty() && result.size() < limit) {
            WeightedTrieNode node = pq.poll();
            // 实际实现中需要同时存储单词字符串,此处简化
            result.add("word_" + node.frequency);
        }
        return result;
    }

    private void collectAllWords(WeightedTrieNode node, String prefix,
                                  PriorityQueue<WeightedTrieNode> pq) {
        if (node.isEndOfWord) pq.offer(node);
        for (Map.Entry<Character, WeightedTrieNode> e : node.children.entrySet()) {
            collectAllWords(e.getValue(), prefix + e.getKey(), pq);
        }
    }

    private WeightedTrieNode findNode(String prefix) {
        WeightedTrieNode current = root;
        for (char ch : prefix.toCharArray()) {
            current = current.children.get(ch);
            if (current == null) return null;
        }
        return current;
    }
}

位运算优化版:适用于仅含小写字母的场景

当字符集仅为26个小写字母时,可用固定大小的数组代替HashMap,显著提升访问速度。

public class ArrayTrieNode {
    ArrayTrieNode[] children = new ArrayTrieNode[26];
    boolean isEndOfWord;

    // 获取子节点,不存在时自动创建
    ArrayTrieNode getOrCreate(char ch) {
        int idx = ch - 'a';
        if (children[idx] == null) {
            children[idx] = new ArrayTrieNode();
        }
        return children[idx];
    }

    ArrayTrieNode get(char ch) {
        return children[ch - 'a'];
    }
}

复杂度分析

操作 时间复杂度 空间复杂度
插入 $O(m)$ $O(m)$
精确查找 $O(m)$ $O(1)$
前缀查找 $O(m)$ $O(1)$
自动补全(k个结果) $O(m + k \cdot L)$ $O(k \cdot L)$

其中 $m$ 为查询字符串长度,$L$ 为平均单词长度。Trie树以空间换时间,在字符串前缀查询密集的场景下具有显著优势。