字典树和后缀数组

MrHe··9 min read

14.2 字典树的应用算法设计

字典树,也叫Trie树或者前缀树,顾名思义,它就是一种专门用来处理字符串前缀的树形结构。每个节点代表一个字符,从根节点到任意一个节点的路径,就构成了一个字符串。它的核心思想就是用空间换时间,利用字符串的公共前缀来降低查询时间的开销。

14.2.1 实现Trie(前缀树)

问题描述

实现一个 Trie (前缀树),包含 insert, search, 和 startsWith 这三个操作。

  • Trie() 初始化前缀树对象。

  • void insert(String word) 向前缀树中插入一个字符串 word

  • boolean search(String word) 如果字符串 word 在前缀树中,返回 true(即,在检索之前已经插入);否则,返回 false

  • boolean startsWith(String prefix) 如果之前已经插入的字符串 word 的前缀之一为 prefix ,返回 true ;否则,返回 false


这个题是Trie的基础,也是面试里最喜欢问的开胃菜。要实现Trie,首先得想明白它的节点长什么样。

一个节点需要包含两个核心信息:

  1. 指向子节点的指针。如果我们的字符集是26个小写英文字母,那就可以用一个大小为26的数组来存,children[0] 代表 'a',children[1] 代表 'b',以此类推。

  2. 一个标记,用来表示从根节点到当前这个节点所构成的路径是不是一个完整的单词。我习惯用一个isEnd的布尔值,或者用一个end计数器,表示以这个节点结尾的单词有多少个。

思路与实现

整个Trie树就由这样的节点构成。我们还需要一个根节点 root,它不代表任何字符,只是作为所有字符串的共同起点。

  • insert(word): 从根节点开始,遍历word的每个字符。对于字符c,我们计算它在children数组中的索引(比如 c - 'a')。

    • 如果 root.children[index]null,说明这个路径不存在,那就创建一个新的TrieNode放进去。

    • 然后,把 root 指针移动到这个子节点上,继续处理下一个字符。

    • word的所有字符都处理完,把最后一个节点的 isEnd 标志设为 true

  • search(word): 和insert非常像,也是从根节点开始顺着word的路径往下走。

    • 如果在路上任何一步发现 root.children[index]null,说明这个单词肯定没插入过,直接返回 false

    • 如果word所有字符都走完了,最后要检查当前节点的 isEnd 是不是 true。是true才说明这是一个完整插入过的单词,否则它可能只是某个更长单词的前缀。

  • startsWith(prefix): 和search几乎一样,唯一的区别是,只要prefix的路径能完整走完,就返回true,我们不关心最后一个节点是不是一个单词的结尾(即不关心isEnd)。

代码实现

class Trie {

    private class TrieNode {
        // isEnd 表示这个节点是否是一个单词的结尾
        boolean isEnd;
        // children 数组,存储指向26个小写字母的子节点
        TrieNode[] children;

        public TrieNode() {
            isEnd = false;
            children = new TrieNode[26];
        }
    }

    private TrieNode root;

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

    // 插入一个单词
    public void insert(String word) {
        TrieNode node = root;
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                node.children[index] = new TrieNode();
            }
            node = node.children[index];
        }
        node.isEnd = true;
    }

    // 搜索一个单词是否存在
    public boolean search(String word) {
        TrieNode node = findNode(word);
        // 不仅节点要存在,而且必须是一个单词的结尾
        return node != null && node.isEnd;
    }

    // 检查是否存在以指定前缀开头的单词
    public boolean startsWith(String prefix) {
        // 只要能找到前缀对应的节点,就说明存在
        return findNode(prefix) != null;
    }

    // 辅助方法:查找一个字符串(或前缀)对应的最后一个节点
    private TrieNode findNode(String str) {
        TrieNode node = root;
        for (char c : str.toCharArray()) {
            int index = c - 'a';
            if (node.children[index] == null) {
                return null;
            }
            node = node.children[index];
        }
        return node;
    }
}

这里我把查找节点的逻辑抽成了一个私有方法 findNode,让searchstartsWith的实现更清晰。对于这个基础题,其实没有所谓的“第二种解法”,因为这就是Trie的标准定义和实现。任何变种都是基于这个核心结构。

14.2.2 最长公共前缀

问题描述

编写一个函数来查找字符串数组中的最长公共前缀。如果不存在公共前缀,返回空字符串 ""

示例: 输入: strs = ["flower","flow","flight"] 输出: "fl"


这个题是个经典老题了,常规解法大家肯定都很熟。

解法一:纵向扫描

这个思路最直观。把第一个字符串"flower"当做基准,然后拿它的每一个字符,去和数组里其他所有字符串的对应位置的字符比较。

  1. 从第一个字符串的第一个字符 f 开始。

  2. 遍历数组中剩下的所有字符串("flow", "flight"),看它们第一个字符是不是也是 f。是的。

  3. 再看第一个字符串的第二个字符 l

  4. 遍历数组中剩下的所有字符串,看它们第二个字符是不是也是 l。是的。

  5. 再看第一个字符串的第三个字符 o

  6. 遍历到"flight"时,发现第三个字符是 i,不匹配了。

  7. 那么最长公共前缀就是 fl,结束。

这个方法的时间复杂度是 O(S),其中 S 是所有字符串的总长度。空间复杂度 O(1)。

class Solution_LCP1 {
    public String longestCommonPrefix(String[] strs) {
        if (strs == null || strs.length == 0) {
            return "";
        }

        // 基准字符串的长度
        int length = strs[0].length();
        // 字符串数组的大小
        int count = strs.length;

        for (int i = 0; i < length; i++) {
            char c = strs[0].charAt(i);
            // 遍历其他字符串
            for (int j = 1; j < count; j++) {
                // 如果其他字符串已经到头了,或者当前字符不匹配
                if (i == strs[j].length() || strs[j].charAt(i) != c) {
                    return strs[0].substring(0, i);
                }
            }
        }
        // 如果循环正常结束,说明第一个字符串本身就是最长公共前缀
        return strs[0];
    }
}

解法二:字典树 (Trie)

既然我们在讲Trie,那自然要用Trie来解。Trie是前缀树,这题问最长公共前缀,简直是量身定做。

思路是什么? 如果所有字符串都有一个公共前缀,那么在Trie树里,它们必然会共享从根节点开始的一条共同路径。当路径出现分叉,或者某个单词在这条路径上结束了,那么公共前缀也就到头了。

  1. 把数组里所有的字符串都插入到Trie树中。

  2. 从根节点 root 开始往下遍历。

  3. 在当前节点,检查它有几个子节点。

    • 如果子节点数量 不等于1,说明从这里开始,路径出现了分叉(比如"flower""flight"l节点后,一个走向o,一个走向i),公共前缀到此为止。

    • 另外,还要检查当前节点本身是不是一个单词的结尾(isEnd == true)。如果是,比如数组是["flow", "flower"],当走到 w 对应的节点时,"flow"已经结束了,那公共前缀最多也就是"flow"

  4. 如果当前节点只有一个子节点,并且它自己也不是单词结尾,说明所有字符串都包含这个路径。我们把当前字符加入结果,然后继续移动到那个唯一的子节点。

  5. 重复这个过程,直到不满足条件。

class Solution_LCP2 {
    class TrieNode {
        // childCount 记录子节点的数量
        int childCount; 
        boolean isEnd;
        TrieNode[] children;

        TrieNode() {
            childCount = 0;
            isEnd = false;
            children = new TrieNode[26];
        }
    }

    public String longestCommonPrefix(String[] strs) {
        if (strs == null || strs.length == 0) {
            return "";
        }

        // 1. 构建Trie树
        TrieNode root = new TrieNode();
        for (String str : strs) {
            // 如果有个空字符串,那公共前缀就是""
            if (str.isEmpty()) {
                return "";
            }
            TrieNode node = root;
            for (char c : str.toCharArray()) {
                int index = c - 'a';
                if (node.children[index] == null) {
                    node.children[index] = new TrieNode();
                    node.childCount++;
                }
                node = node.children[index];
            }
            node.isEnd = true;
        }

        // 2. 在Trie树中查找最长公共前缀
        StringBuilder lcp = new StringBuilder();
        TrieNode node = root;
        while (true) {
            // 如果节点出现分叉,或者一个单词在这里结束了(但strs.length > 1)
            // 注意:如果只有一个单词,isEnd不应中断
            if (node.childCount != 1 || (node.isEnd && strs.length > 1)) {
                break;
            }

            // 找到那个唯一的子节点
            int nextIndex = -1;
            for (int i = 0; i < 26; i++) {
                if (node.children[i] != null) {
                    nextIndex = i;
                    break;
                }
            }

            // 理论上 childCount==1 就一定能找到
            if(nextIndex != -1) {
                lcp.append((char) ('a' + nextIndex));
                node = node.children[nextIndex];
            } else {
                 // 正常情况下不会走到这里
                break;
            }
        }
        return lcp.toString();
    }
}

对比一下,解法一简单粗暴,容易想到。解法二利用了数据结构的优势,虽然代码量大了些,但思路更贴合“前缀”这个概念。在某些场景下,如果Trie树是预先构建好的,那么查询最长公共前缀的效率会非常高。

14.2.3 单词替换

问题描述

在英语中,我们有一个叫做 词根 (root) 的概念,可以词根后面添加其他一些词组成另一个较长的单词——我们称这个词为 继承词 (successor)。例如,词根an,跟随着单词 other,可以形成新的单词 another

现给定一个由许多词根组成的词典 dictionary 和一个用空格分隔单词的句子 sentence。你需要将句子中的所有继承词用最短的词根来替换掉。如果继承词有许多可以形成它的词根,则用最短的那个词根替换它。

如果继承词没有对应的词根,则不要替换该单词。

示例: 输入: dictionary = ["cat","bat","rat"], sentence = "the cattle was rattled by the battery" 输出: "the cat was rat by the bat"


这个题目的本质就是:对于句子里的每个单词,找到它在词典里的最短前缀。

解法一:哈希集合 + 暴力前缀匹配

一个很直接的想法是,把所有词根都存到一个哈希集合(HashSet)里,这样查询一个词根是否存在就是 O(1) 的时间复杂度(平均情况)。

然后,我们把句子按空格切分成单词数组。对每个单词,我们从短到长遍历它的所有前缀:

  1. 对于单词 cattle:

    • 检查前缀 c 是否在哈希集合里?否。

    • 检查前缀 ca 是否在哈希集合里?否。

    • 检查前缀 cat 是否在哈希集合里?是!找到了,而且因为我们是从短到长遍历的,所以这一定是能匹配上的最短词根。直接用 cat 替换 cattle,然后处理下一个单词。

  2. 对于单词 rattled:

    • r? 否。 ra? 否。 rat? 是!用 rat 替换。
  3. 如果一个单词的所有前缀都不在词典里,那就不替换。

import java.util.HashSet;
import java.util.List;
import java.util.Set;

class Solution_WR1 {
    public String replaceWords(List<String> dictionary, String sentence) {
        Set<String> dictSet = new HashSet<>(dictionary);
        String[] words = sentence.split(" ");
        StringBuilder result = new StringBuilder();

        for (int i = 0; i < words.length; i++) {
            String word = words[i];
            String prefix = "";
            for (int j = 1; j <= word.length(); j++) {
                prefix = word.substring(0, j);
                if (dictSet.contains(prefix)) {
                    // 找到了最短的词根前缀,跳出内层循环
                    break;
                }
            }
            // 如果循环结束prefix还是原来的word.substring(...)的结果,说明找到了
            // 如果内层循环是因为 j > word.length() 跳出的, 说明没找到, prefix就是word本身
            // 这里逻辑有点绕,更清晰的写法是设定一个标志位

            // 更好的写法
            String replacement = word;
            for (int j = 1; j <= word.length(); j++) {
                String currentPrefix = word.substring(0, j);
                if (dictSet.contains(currentPrefix)) {
                    replacement = currentPrefix;
                    break;
                }
            }

            result.append(replacement);
            if (i < words.length - 1) {
                result.append(" ");
            }
        }
        return result.toString();
    }
}

这个方法的时间复杂度是多少?假设句子有 N 个单词,每个单词平均长度为 L,词典大小为 M。构建哈希集合是 O(M * L_dict)。对每个单词,我们最多要进行 L 次substringcontains操作,substring是 O(L),所以是 O(L^2)。总的就是 O(N * L^2)。

解法二:字典树 (Trie)

一看这题又是“前缀”,Trie 的警觉立马就要提起来。用Trie怎么做?

  1. 把词典里所有的词根 dictionary全部插入到Trie树中。

  2. 把句子 sentence 按空格切分成单词数组。

  3. 遍历每个单词,比如 cattle

  4. 从Trie的根节点开始,顺着 c-a-t-t-l-e 的路径往下走。

  5. 在走的每一步,都要检查当前节点是不是一个词根的结尾(isEnd == true)。

    • 走到 c,不是结尾。

    • 走到 a,不是结尾。

    • 走到 t,发现这个节点 isEndtrue!这说明 cat 是一个词根。因为我们是从根节点一步步往下走的,所以第一个碰到的 isEnd 节点,对应的路径一定就是最短的词根前缀。

    • 此时,我们就可以停止对 cattle 的搜索,直接用 cat 作为替换词。

  6. 如果整个单词都走完了,路径上也没碰到任何一个 isEndtrue 的节点,说明它没有词根前缀,保留原单词。

这个过程,对每个单词的搜索,只需要从头到尾遍历一遍,时间复杂度是 O(L),比解法一的 O(L^2) 要高效。

import java.util.List;

class Solution_WR2 {
    class TrieNode {
        boolean isEnd;
        TrieNode[] children = new TrieNode[26];
    }

    public String replaceWords(List<String> dictionary, String sentence) {
        TrieNode root = new TrieNode();
        // 1. 构建词根Trie树
        for (String rootWord : dictionary) {
            TrieNode node = root;
            for (char c : rootWord.toCharArray()) {
                int index = c - 'a';
                if (node.children[index] == null) {
                    node.children[index] = new TrieNode();
                }
                node = node.children[index];
            }
            node.isEnd = true;
        }

        String[] words = sentence.split(" ");
        StringBuilder result = new StringBuilder();

        // 2. 遍历句子中的每个单词,在Trie中查找最短前缀
        for (int i = 0; i < words.length; i++) {
            String word = words[i];
            String replacement = findShortestRoot(word, root);
            result.append(replacement);
            if (i < words.length - 1) {
                result.append(" ");
            }
        }
        return result.toString();
    }

    private String findShortestRoot(String word, TrieNode root) {
        TrieNode node = root;
        StringBuilder prefix = new StringBuilder();
        for (char c : word.toCharArray()) {
            int index = c - 'a';
            // 如果路径断了, 或者已经找到了一个词根, 就可以停了
            if (node.children[index] == null) {
                // 没找到任何词根
                return word;
            }

            prefix.append(c);
            node = node.children[index];

            if (node.isEnd) {
                // 找到了最短的词根
                return prefix.toString();
            }
        }
        // 整个单词都匹配完了也没找到词根
        return word;
    }
}

Trie的解法在处理这类前缀匹配问题上,效率优势体现得淋漓尽致。

14.3 后缀数组的应用算法设计

聊完了Trie,我们进入另一个大杀器:后缀数组。后缀数组(Suffix Array, SA)主要是为了解决子串(substring)相关的问题。它和它的好搭档LCP数组(Longest Common Prefix Array)结合,威力无穷。

简单说一下核心概念:

  • 后缀:字符串 S = "banana",它的所有后缀是 "banana", "anana", "nana", "ana", "na", "a"

  • 后缀数组 sa:把所有后缀按字典序排序后,sa[i] 存储的是排在第i位的后缀在原字符串中的起始下标。

  • 名次数组 rankrank[i] 存储的是以i为起始下标的后缀在排序后的排名。sarank是互逆的。

  • LCP数组 heightheight[i] 存储的是排序后第i个后缀(即sa[i])和第i-1个后缀(即sa[i-1])的最长公共前缀(LCP)的长度。height[0]通常定义为0。

后缀数组的构造算法(如DC3, SA-IS)比较复杂,在面试中通常不要求手写,我们这里就假设已经有现成的工具可以生成saheight数组,重点关注如何利用它们来解题。

14.3.1 字符串的不同子串个数

问题描述

给定一个字符串 S,求 S 的所有不同子串的个数。

示例: 输入: s = "ababa" 输出: 9 解释: 不同的子串是 "a", "b", "ab", "ba", "aba", "bab", "baba", "ababa"。 a 有两个, b有两个,ab有两个,aba有两个,我们只算一次。"a", "b", "a", "b", "a" -> "a","b". 子串 "ab", "ba", "ab", "a". "aba", "bab", "ba". "ababa"


解法一:哈希集合

最无脑的办法,生成所有子串,然后扔到 HashSet 里,最后返回 set.size()

  1. 两层循环,外层i从0到n-1,内层jin-1

  2. s.substring(i, j+1) 就是一个子串。

  3. 把它加入 HashSet

这个方法简单,但如果字符串很长,生成的子串数量是 O(N^2) 级别,每个子串的长度可能是O(N),存入哈希表也要时间和空间,整体效率不高,容易内存溢出。

import java.util.HashSet;
import java.util.Set;

class Solution_DS1 {
    public int countDistinctSubstrings(String s) {
        Set<String> distinctSubstrings = new HashSet<>();
        int n = s.length();
        for (int i = 0; i < n; i++) {
            for (int j = i; j < n; j++) {
                distinctSubstrings.add(s.substring(i, j + 1));
            }
        }
        return distinctSubstrings.size();
    }
}

解法二:后缀数组

这是后缀数组的经典应用。核心思想是:任何一个子串,都必然是某个后缀的前缀

比如 s = "ababa",子串 "bab" 就是后缀 "baba" (下标为1) 的前缀。

那么,问题就转化为:求所有后缀的所有不同前缀的个数

  1. 我们先把所有后缀列出来,排好序。

  2. 如果不考虑重复,一个长度为 k 的后缀,它能提供 k 个前缀。所以所有后缀能提供的总前缀数(包含重复)是 (n) + (n-1) + ... + 1 = n * (n+1) / 2

  3. 重复的前缀是怎么产生的?当两个后缀有公共前缀时,这些公共前缀就被重复计算了。比如后缀 "ababa""aba",它们的公共前缀是 "aba"。那么 "a", "ab", "aba" 这三个前缀就被计算了两次。

  4. LCP数组 height 恰好告诉我们,排序后相邻两个后缀的最长公共前缀的长度。height[i] 就是 suffix(sa[i])suffix(sa[i-1]) 的LCP长度。这意味着,它们有 height[i] 个公共的前缀,这些前缀在计算 suffix(sa[i-1]) 的前缀时已经算过了,所以在计算 suffix(sa[i]) 的前缀时,这 height[i] 个就是重复的。

  5. 所以,总的不同子串个数 = 所有后缀的前缀总数 - 因重复而多计算的个数。

    • 总数 = n * (n+1) / 2

    • 重复数 = sum(height[i]) for i from 1 to n-1.

    • 结果 = n * (n+1) / 2 - sum(height)

这个公式是解决这类问题的关键。我们只需要一个构造后缀数组和LCP数组的工具类即可。

// 假设我们有一个 SuffixArrayUtil 类,能生成 sa 和 height 数组
// 在实际工程或竞赛中,这通常是一个预先写好的模板
class Solution_DS2 {
    public int countDistinctSubstrings(String s) {
        int n = s.length();
        // 实际使用时, 你需要一个后缀数组的实现
        int[] sa = SuffixArrayUtil.getSa(s);
        int[] height = SuffixArrayUtil.getHeight(s, sa);

        long totalSubstrings = (long) n * (n + 1) / 2;
        long repeatedSubstrings = 0;
        for (int h : height) {
            repeatedSubstrings += h;
        }

        return (int) (totalSubstrings - repeatedSubstrings);
    }
}

// SuffixArrayUtil 的一个简单(但效率不高O(N^2logN))的实现,仅作演示
// 竞赛级实现通常用DC3或SA-IS, O(N)或O(NlogN)
class SuffixArrayUtil {
    public static int[] getSa(String s) {
        // ... 复杂的实现 ...
        return new int[s.length()]; // 占位
    }
    public static int[] getHeight(String s, int[] sa) {
        // ... 复杂的实现 ...
        return new int[s.length()]; // 占位
    }
}

14.3.2 最长重复子串

问题描述

给定一个字符串 S,找出其中最长的重复子串。如果有多个答案,返回任意一个即可。如果不存在重复子串,返回空字符串。

示例: 输入: s = "banana" 输出: "ana" (或 "nan")


解法一:暴力枚举 + 哈希

这个思路和上一题的暴力解类似。我们可以二分答案,猜最长重复子串的长度 k

  1. 二分查找长度 k,范围从 1 到 n-1。

  2. 对于一个给定的 k,我们要检查是否存在长度为 k 的重复子串。

  3. 怎么检查?生成所有长度为 k 的子串,把它们放到 HashSet 里。如果某个子串在 add 的时候失败了(因为已经存在),说明找到了长度为 k 的重复子串,那么我们可以尝试更长的 kleft = mid + 1)。

  4. 如果所有长度为 k 的子串都成功加入了 set,说明没有这个长度的重复子串,我们要尝试短一点的 kright = mid - 1)。

这个方法结合了二分和哈希,比纯暴力要好,但复杂度依然不低。检查一个k的时间是 O(Nk),二分是 logN,总共是 O(Nk*logN)。

import java.util.HashSet;
import java.util.Set;

class Solution_LRS1 {
    public String longestDupSubstring(String s) {
        int n = s.length();
        int left = 1, right = n - 1;
        int start = 0, maxLen = 0;

        while (left <= right) {
            int mid = left + (right - left) / 2;
            int foundStart = search(s, mid);
            if (foundStart != -1) {
                // 找到了长度为 mid 的重复子串
                maxLen = mid;
                start = foundStart;
                left = mid + 1; // 尝试更长的
            } else {
                right = mid - 1; // 只能尝试更短的
            }
        }
        return maxLen > 0 ? s.substring(start, start + maxLen) : "";
    }

    // 检查是否存在长度为 len 的重复子串
    private int search(String s, int len) {
        Set<String> seen = new HashSet<>();
        for (int i = 0; i <= s.length() - len; i++) {
            String sub = s.substring(i, i + len);
            if (!seen.add(sub)) {
                return i; // 找到了重复的
            }
        }
        return -1;
    }
}

注意:这个解法在处理大字符串时,substringHashSet可能会因为大量对象创建而导致性能问题或内存超限,所以会使用滚动哈希(Rabin-Karp)进行优化,这里就不展开了。

解法二:后缀数组

又是后缀数组的主场。

思路转换: 一个重复出现的子串,它必然是至少两个不同后缀的公共前缀。

比如 s = "banana",重复子串 "ana",它既是后缀 "anana" 的前缀,也是后缀 "ana" 的前缀。

我们的目标是找最长的重复子串,这就等价于找所有后缀对的最长公共前缀(LCP)的最大值

这里有一个非常重要的性质: 任意两个后缀 suffix(i)suffix(j) 的LCP,等于它们在排好序的后缀数组中,中间所有相邻后缀对的LCP的最小值。即 LCP(i, j) = min(height[k]),其中 krank[i]+1rank[j] 之间。

这个性质告诉我们,我们不需要去计算任意两个后缀的LCP,我们只要求出 height 数组。height 数组中的最大值,就是所有相邻后缀对 LCP 的最大值,也自然是所有后缀对LCP的最大值。

所以,解法极其简洁:

  1. 计算字符串 s 的后缀数组 sa

  2. 基于 sa 计算出LCP数组 height

  3. 找到 height 数组中的最大值。这个值就是最长重复子串的长度。

// 同样,假设我们有 SuffixArrayUtil 工具类
class Solution_LRS2 {
    public String longestDupSubstring(String s) {
        int n = s.length();
        if (n == 0) return "";

        int[] sa = SuffixArrayUtil.getSa(s);
        int[] height = SuffixArrayUtil.getHeight(s, sa);

        int maxLen = 0;
        int startIndex = -1;

        for (int i = 0; i < n; i++) {
            if (height[i] > maxLen) {
                maxLen = height[i];
                // 这个LCP是 suffix(sa[i]) 和 suffix(sa[i-1]) 的
                // 我们取 sa[i] 作为起始点即可
                startIndex = sa[i];
            }
        }

        if (startIndex == -1) {
            return "";
        }
        return s.substring(startIndex, startIndex + maxLen);
    }
}