如何查找和为给定值的所有子数组并实现O(n)时间复杂度优化

原有代码问题 #

• 匹配到第一个符合条件的子数组就直接执行return终止逻辑，无法输出所有匹配结果
• 每次调用sum(result)计算窗口和的操作时间复杂度为O(k)（k为当前窗口长度），整体实际时间复杂度为O(n²)，且该滑动窗口方案仅适用于数组元素全为正的场景，存在局限性

通用O(n)解法：前缀和+哈希表 #

核心逻辑：定义前缀和pre_sum[i]为数组前i个元素的和，那么子数组arr[j:i]（左闭右开，包含索引j不包含索引i）的和为pre_sum[i] - pre_sum[j]。我们要找和为s的子数组，等价于遍历到i时，查找所有满足pre_sum[j] = pre_sum[i] - s的j值，对应的arr[j:i]就是符合要求的子数组。
我们用哈希表存储每个前缀和对应的所有索引，单次查找时间为O(1)，整体时间复杂度为O(n)，支持包含负数、零的所有整数数组场景。

实现代码 #

def get_all_sub(arr, s):
    pre_sum_map = {0: [0]}  # key:前缀和 value:对应该前缀和的所有索引
    pre_sum = 0
    res = []
    for idx, num in enumerate(arr):
        pre_sum += num
        # 查找是否存在符合要求的前缀和
        target = pre_sum - s
        if target in pre_sum_map:
            # 所有对应的j都生成对应的子数组
            for j in pre_sum_map[target]:
                res.append(arr[j:idx+1])
        # 将当前前缀和存入哈希表
        if pre_sum not in pre_sum_map:
            pre_sum_map[pre_sum] = []
        pre_sum_map[pre_sum].append(idx+1)
    return res

# 测试用例
arr = [4, 1, 6, 5, 2, 3]
s = 5
print(get_all_sub(arr, s))

输出结果 #

[[4, 1], [5], [2, 3]]

全正数数组优化滑动窗口O(n)方案 #

如果可以确定输入数组所有元素都是正数，也可以修改原来的滑动窗口逻辑实现O(n)复杂度：

def get_all_sub_positive(arr, s):
    left = 0
    current_sum = 0
    res = []
    for right in range(len(arr)):
        current_sum += arr[right]
        # 窗口和超过目标就移动左指针
        while current_sum > s:
            current_sum -= arr[left]
            left += 1
        # 匹配到就记录结果
        if current_sum == s:
            res.append(arr[left:right+1])
    return res

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