如何优化计算和等于指定值的子串数量的超时Python代码

你原有的代码时间复杂度为O(n²)，每轮循环都需要生成新的累加数组并遍历计数，输入字符串长度较大时会触发超时。我们可以采用前缀和+哈希表计数的方法，将时间复杂度降至O(n)，可轻松应对超长输入场景。

核心原理 #

定义前缀和pre_sum[i]为序列前i个元素的和，我们需要统计的连续子串和等于m的数量，等价于统计满足pre_sum[j] - pre_sum[i] = m的(i,j)配对总数，也就是遍历到第j位时，值为pre_sum[j] - m的前缀和已经出现过的次数。我们用哈希表实时记录每个前缀和的出现次数，单次遍历即可完成统计。

优化后代码 #

from collections import defaultdict

m = 9
numbers = list(map(int, S.split(',')))
# 记录每个前缀和的出现次数，初始前缀和0对应子串从首位开始的场景
prefix_cnt = defaultdict(int)
prefix_cnt[0] = 1
current_sum = 0
result = 0

for num in numbers:
    current_sum += num
    # 累加符合要求的历史前缀和数量
    result += prefix_cnt.get(current_sum - m, 0)
    # 更新当前前缀和的计数
    prefix_cnt[current_sum] += 1

print(result)

优化收益 #

• 时间复杂度从O(n²)降到O(n)，即使是长度10万级别的输入也可以在1秒内完成计算
• 空间复杂度从原方案的O(n)降到O(k)，k为不同前缀和的数量，多数场景下远小于序列长度n

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