前缀树（Trie）

MrHe··4 min read

前缀树是一种专门处理字符串匹配的树形数据结构，又叫字典树。它的核心思想是空间换时间，利用字符串的公共前缀来减少查询时间，从而提高效率。

前缀树的核心思想

前缀树的每个节点都代表一个字符，从根节点到某一节点的路径连接起来就是一个字符串。通过存储字符之间的层级关系，前缀树可以高效地进行字符串的插入、查找和删除操作。

经典题型：实现一个前缀树

实现一个 Trie 类，包含 insert、search 和 startsWith 三个方法。

方法一：基础数组实现

使用数组来存储子节点，是最直观的实现方式。

class TrieNode {
    // 子节点，这里假设只包含小写字母
    private TrieNode[] children;
    // 标记是否是单词的结尾
    private boolean isEnd;

    public TrieNode() {
        children = new TrieNode[26]; // 26个小写字母
        isEnd = false;
    }

    public boolean containsKey(char ch) {
        return children[ch - 'a'] != null;
    }

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

    public void put(char ch, TrieNode node) {
        children[ch - 'a'] = node;
    }

    public void setEnd() {
        isEnd = true;
    }

    public boolean isEnd() {
        return isEnd;
    }
}

class Trie {
    private TrieNode root;

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

    // 插入一个单词
    public void insert(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (!node.containsKey(ch)) {
                node.put(ch, new TrieNode());
            }
            node = node.get(ch);
        }
        node.setEnd();
    }

    // 查找一个单词是否存在
    public boolean search(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (!node.containsKey(ch)) {
                return false;
            }
            node = node.get(ch);
        }
        return node.isEnd();
    }

    // 查找是否有以该前缀开头的单词
    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (int i = 0; i < prefix.length(); i++) {
            char ch = prefix.charAt(i);
            if (!node.containsKey(ch)) {
                return false;
            }
            node = node.get(ch);
        }
        return true;
    }
}

这种方法的时间和空间复杂度：

  • 插入：O(m)，其中 m 是单词的长度

  • 查找：O(m)

  • 前缀查找：O(m)

  • 空间：O(26^n)，其中 n 是树的深度

方法二：HashMap实现

使用HashMap来存储子节点，可以支持更广泛的字符集，而不仅仅是小写字母。

class TrieNode {
    Map<Character, TrieNode> children;
    boolean isEnd;

    public TrieNode() {
        children = new HashMap<>();
        isEnd = false;
    }
}

class Trie {
    private TrieNode root;

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

    public void insert(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (!node.children.containsKey(ch)) {
                node.children.put(ch, new TrieNode());
            }
            node = node.children.get(ch);
        }
        node.isEnd = true;
    }

    public boolean search(String word) {
        TrieNode node = root;
        for (int i = 0; i < word.length(); i++) {
            char ch = word.charAt(i);
            if (!node.children.containsKey(ch)) {
                return false;
            }
            node = node.children.get(ch);
        }
        return node.isEnd;
    }

    public boolean startsWith(String prefix) {
        TrieNode node = root;
        for (int i = 0; i < prefix.length(); i++) {
            char ch = prefix.charAt(i);
            if (!node.children.containsKey(ch)) {
                return false;
            }
            node = node.children.get(ch);
        }
        return true;
    }
}

方法三：压缩前缀树（Radix Tree）

压缩前缀树是前缀树的一种优化，它将只有一个子节点的链进行压缩，从而减少空间使用。

class RadixNode {
    Map<String, RadixNode> children;
    boolean isEnd;

    public RadixNode() {
        children = new HashMap<>();
        isEnd = false;
    }
}

class RadixTree {
    private RadixNode root;

    public RadixTree() {
        root = new RadixNode();
    }

    public void insert(String word) {
        RadixNode node = root;
        int i = 0;
        while (i < word.length()) {
            boolean found = false;
            for (Map.Entry<String, RadixNode> entry : node.children.entrySet()) {
                String edge = entry.getKey();
                RadixNode child = entry.getValue();

                // 找到共同前缀
                int j = 0;
                while (j < edge.length() && i + j < word.length() && edge.charAt(j) == word.charAt(i + j)) {
                    j++;
                }

                if (j > 0) {
                    found = true;

                    if (j == edge.length() && j == word.length() - i) {
                        // 完全匹配
                        child.isEnd = true;
                        return;
                    } else if (j == edge.length()) {
                        // 沿着当前边继续
                        node = child;
                        i += j;
                        break;
                    } else if (j < edge.length()) {
                        // 需要分裂当前边
                        String prefix = edge.substring(0, j);
                        String suffix = edge.substring(j);

                        // 创建新节点
                        RadixNode newNode = new RadixNode();
                        newNode.children.put(suffix, child);

                        // 更新当前节点的子节点
                        node.children.remove(edge);
                        node.children.put(prefix, newNode);

                        // 处理剩余字符
                        if (j < word.length() - i) {
                            String remaining = word.substring(i + j);
                            RadixNode endNode = new RadixNode();
                            endNode.isEnd = true;
                            newNode.children.put(remaining, endNode);
                        } else {
                            newNode.isEnd = true;
                        }

                        return;
                    }
                }
            }

            if (!found) {
                // 没有匹配的边，直接添加
                node.children.put(word.substring(i), new RadixNode());
                node.children.get(word.substring(i)).isEnd = true;
                return;
            }
        }
    }

    public boolean search(String word) {
        RadixNode node = root;
        int i = 0;
        while (i < word.length()) {
            boolean found = false;
            for (Map.Entry<String, RadixNode> entry : node.children.entrySet()) {
                String edge = entry.getKey();

                if (i + edge.length() <= word.length() && word.substring(i, i + edge.length()).equals(edge)) {
                    i += edge.length();
                    node = entry.getValue();
                    found = true;
                    break;
                }
            }

            if (!found) {
                return false;
            }
        }

        return node.isEnd;
    }

    public boolean startsWith(String prefix) {
        RadixNode node = root;
        int i = 0;
        while (i < prefix.length()) {
            boolean found = false;
            for (Map.Entry<String, RadixNode> entry : node.children.entrySet()) {
                String edge = entry.getKey();

                if (i + edge.length() <= prefix.length() && prefix.substring(i, i + edge.length()).equals(edge)) {
                    i += edge.length();
                    node = entry.getValue();
                    found = true;
                    break;
                }
            }

            if (!found) {
                return false;
            }
        }

        return true;
    }
}

压缩前缀树在最坏情况下的空间复杂度是 O(n)，其中 n 是所有字符串的总长度，而在最佳情况下（字符串有大量公共前缀），空间复杂度可以接近 O(m)，其中 m 是不同字符串的数量。

前缀树的应用场景

前缀树在以下场景中有广泛应用：

  1. 字符串检索：快速判断一个字符串是否在集合中

  2. 前缀匹配：找出所有具有特定前缀的字符串

  3. 自动补全：根据用户输入的前缀，提示可能的完整单词

  4. IP路由：最长前缀匹配

  5. 拼写检查：检查单词拼写是否正确

  6. 词频统计：统计文本中单词出现的频率

前缀树的解题模板

// 基本前缀树实现
class TrieTemplate {
    private TrieNode root;

    private static class TrieNode {
        Map<Character, TrieNode> children;
        boolean isEnd;

        public TrieNode() {
            children = new HashMap<>();
            isEnd = false;
        }
    }

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

    public void insert(String word) {
        TrieNode node = root;
        for (char ch : word.toCharArray()) {
            if (!node.children.containsKey(ch)) {
                node.children.put(ch, new TrieNode());
            }
            node = node.children.get(ch);
        }
        node.isEnd = true;
    }

    public boolean search(String word) {
        TrieNode node = searchPrefix(word);
        return node != null && node.isEnd;
    }

    public boolean startsWith(String prefix) {
        return searchPrefix(prefix) != null;
    }

    private TrieNode searchPrefix(String prefix) {
        TrieNode node = root;
        for (char ch : prefix.toCharArray()) {
            if (!node.children.containsKey(ch)) {
                return null;
            }
            node = node.children.get(ch);
        }
        return node;
    }
}

前缀树是字符串处理中的重要数据结构，掌握了前缀树的实现和应用，可以在很多字符串相关问题中找到高效的解决方案。在实际工程中，前缀树的变体如压缩前缀树、后缀树等也有广泛应用。