求更优算法：统计字符串列表中的后缀匹配对数量

问题描述 #

统计字符串列表中满足以下条件的i<j配对总数：words[i]是words[j]的后缀（例如，"door"是"backdoor"的后缀，"a"是"cba"的后缀）。

现有暴力实现 #

版本1（较慢） #

通过反转字符串并排序，遍历每个字符串统计后续以它为前缀的反转字符串数量，时间复杂度O(n log n + nL)（L为字符串平均长度）：

def count_same_endings(words):
    """
    Count the number of word pairs with the same ending.

    Parameters:
    words (list): List of words.

    Returns:
    int: The count of word pairs with the same ending.
    """
    strings = sorted([word[::-1] for word in words])
    count = 0
    i = 0
    while i < len(words):
        j = 1
        while i + j < len(words) and strings[i + j].startswith(strings[i]):
            count += 1
            j += 1
        i += 1
    return count

words = ['cba', 'a', 'a', 'a', 'a', 'b', 'ba', 'ca']
print(count_same_endings(words))  # 输出19

版本2（较快） #

通过统计所有后缀片段的出现次数来计数，时间复杂度O(nL^2)（L为字符串最大长度）：

def count_end_slices(strings):
    slice_count = {}
    strings_count = {}
    count = 0

    # Iterate over each word in the list
    for string in strings:
        length = len(string)

        # Check if this word is already the suffix of some previous words
        if string in slice_count:
            count += slice_count[string]

        # Generate slices that include the end of the string
        for start in range(length):
            slice = string[start:]

            # Count the occurrences of each slice
            if slice in slice_count:
                # Check if previous words are the suffix of this string
                if slice in strings_count:
                    count += strings_count[slice] if slice != string else 0
                slice_count[slice] += 1
            else:
                slice_count[slice] = 1

        # Account for this word
        if string in strings_count:
            strings_count[string] += 1
        else:
            strings_count[string] = 1

    return slice_count, count

优化方案：字典树（前缀树） #

利用反转字符串将后缀匹配转化为前缀匹配，结合字典树实现线性时间复杂度（相对于所有字符串的总长度）的统计：

核心思路 #

1. 反转所有字符串，将"判断A是B的后缀"转化为"判断反转后的A是反转后的B的前缀"；
2. 按字符串长度从小到大排序，保证短字符串先被插入字典树（短字符串才可能是长字符串的前缀）；
3. 用字典树记录每个前缀的结束次数，遍历每个字符串时累加所有匹配前缀的结束次数，再将当前字符串插入字典树。

实现代码 #

class TrieNode:
    def __init__(self):
        self.children = {}
        self.end_count = 0  # 记录以当前前缀结尾的字符串数量

def count_suffix_pairs(words):
    # 反转所有字符串
    reversed_words = [word[::-1] for word in words]
    # 按字符串长度从小到大排序，保证短字符串先处理
    reversed_words.sort(key=lambda x: len(x))
    
    root = TrieNode()
    count = 0
    
    for s in reversed_words:
        current = root
        # 遍历当前字符串的所有前缀，累加匹配的数量
        for c in s:
            if c not in current.children:
                break  # 后续前缀不存在，停止遍历
            current = current.children[c]
            count += current.end_count
        
        # 将当前字符串插入字典树
        current = root
        for c in s:
            if c not in current.children:
                current.children[c] = TrieNode()
            current = current.children[c]
        current.end_count += 1
    
    return count

# 测试示例
words = ['cba', 'a', 'a', 'a', 'a', 'b', 'ba', 'ca']
print(count_suffix_pairs(words))  # 输出19

复杂度分析 #

• 时间复杂度：O(T)，其中T是所有字符串的总长度，每个字符最多被遍历两次（一次查询，一次插入）；
• 空间复杂度：O(T)，字典树的总节点数不超过所有字符串的总长度。

内容的提问来源于stack exchange，提问作者Specter43