146字符串首尾相同子序列计数

题解

小 M 拿到了一个仅由小写字母组成的字符串，她想知道在这个字符串中，有多少个子序列的首尾字符相同。子序列的定义是：从原字符串中按原顺序取出若干字符（可以不连续）组成的新字符串。

例如，对于字符串 “arcaea”，其子序列包括 “aca”, “ara”, “aaa” 等，这些子序列的首尾字符都是相同的。

你需要计算满足这一条件的子序列数量，并输出对 998244353 取模的结果。

解题思路

1. 统计字符出现位置：
   使用 defaultdict 来存储每个字符在字符串中出现的所有位置。

2. 计算子序列数量：
   对于每个字符，考虑其所有可能的首尾位置组合 (i, j)，其中 $i < j$。对于每一对 (i, j)，位于 $i$ 和 $j$ 之间的字符可以选择或不选择，因此有 $2^{j - i - 1}$ 种可能的子序列。

3. 累加结果：
   将所有符合条件的子序列数量累加，同时加上单个字符的情况（每个字符本身也是一个有效的子序列）。

4. 取模处理：
   最终结果对 998244353 取模。

代码实现

from collections import defaultdict

def solution(s: str) -> int:
    d = defaultdict(list)
    for i, c in enumerate(s):
        d[c].append(i)
    res = 0
    mod = 998244353
    for k in d:
        v = d[k]
        for i in range(len(v)):
            for j in range(i + 1, len(v)):
                res += pow(2, v[j] - v[i] - 1, mod)
    return (res + len(s)) % mod

if __name__ == '__main__':
    print(solution("arcaea") == 28)
    print(solution("abcabc") == 18)
    print(solution("aaaaa") == 31)

复杂度分析

• 时间复杂度：O(N^2)，其中 N 是字符串的长度。主要消耗在于双重循环遍历每个字符的所有位置组合。
• 空间复杂度：O(N)，用于存储每个字符出现的位置。

注意事项

• 模数应为 998244353，请确保代码中 mod 的值正确。
• 由于可能存在大量的子序列，务必在计算中进行取模操作以防止整数溢出。