求长度不超过k且和等于s的连续子数组数量算法优化

问题描述 #

给定整数数组arr、正整数k、整数s，要求统计所有满足条件的非空连续子数组数量，条件为子数组长度不大于k且子数组元素和等于s。

示例说明 #

输入：arr = [1, 2, 4, -1, 6, 1]，k = 3，s = 6
输出：solution(arr, k, s) = 3
符合条件的子数组如下：

• 长度为1的子数组：[6]，共1个
• 长度为2的子数组：[2, 4]，共1个
• 长度为3的子数组：[-1, 6, 1]，共1个

注意：子数组[1, 2, 4, -1]虽然和为6，但长度超过k，不符合要求，因此不计入结果。

原代码超时原因 #

你当前的实现用字典存储每个前缀和对应的所有索引列表，每次查询到目标前缀和时需要线性遍历所有索引判断长度限制，最坏时间复杂度为O(n²)，当数组长度较大时就会触发超时。

优化实现思路 #

因为前缀和的索引是按遍历顺序递增存入列表的，天然有序，因此可以用二分查找替代线性遍历，统计满足idx >= i - k的索引数量，单次查询时间复杂度降为O(logm)（m为对应前缀和的出现次数），整体时间复杂度优化为O(n logn)，可通过大数据量测试用例。

优化后代码 #

import collections
import bisect

def solution(arr, k, s):
    res = 0
    curr = 0
    # 键为前缀和，值为前缀和对应索引的有序列表
    prefix_map = collections.defaultdict(list)
    prefix_map[0].append(-1)
    
    for i, num in enumerate(arr):
        curr += num
        target = curr - s
        if target in prefix_map:
            idx_list = prefix_map[target]
            # 二分查找第一个大于等于 i - k 的索引位置
            left_pos = bisect.bisect_left(idx_list, i - k)
            # 满足条件的数量为列表总长度减去左边界位置
            res += len(idx_list) - left_pos
        # 将当前前缀和对应的索引存入列表，保持有序性
        prefix_map[curr].append(i)
    return res

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