Leetcode 剑指 Offer II 010. 和为 k 的子数组

news2024/11/18 19:50:10

题目难度: 中等

原题链接

今天继续更新 Leetcode 的剑指 Offer（专项突击版）系列, 大家在公众号 算法精选 里回复 剑指offer2 就能看到该系列当前连载的所有文章了, 记得关注哦~

题目描述

给定一个整数数组和一个整数 k ，请找到该数组中和为 k 的连续子数组的个数。

示例 1：

输入:nums = [1,1,1], k = 2
输出: 2
解释: 此题 [1,1] 与 [1,1] 为两种不同的情况

示例 2：

输入:nums = [1,2,3], k = 3
输出: 2

提示:

1 <= nums.length <= 2 * 10^4
-1000 <= nums[i] <= 1000
-10^7 <= k <= 10^7

题目思考

如何尽可能优化时间复杂度?

解决方案

思路

分析题目, 最容易想到的做法就是暴力两重循环: 外层循环固定每个下标起点, 内层循环往后遍历, 累加当前的子数组和, 等于 k 时计入最终结果中
但这种做法时间效率太低 (O(N^2)), 如何优化呢?
基于暴力法进行分析, 我们对于每个起点都遍历了一遍从它开始的子数组和, 但是这些子数组和完全可以通过从 0 开始的两个前缀和相减所得
例如要求下标[2,3]子数组的和, 可以用下标[0,3]子数组的和减去下标[0,1]的子数组和即可得到
这样只需要遍历一遍数组, 存下所有的前缀和, 就能得到任意连续子数组的和了
所以具体思路如下:
1. 使用一个变量 sm 记录当前前缀和
2. 使用一个计数字典 cnts 统计当前前缀和的出现次数, 注意初始时没有元素, 前缀和为 0, 所以它的初始计数为 1
3. 在遍历每个元素时, 先更新前缀和 sm
4. 然后将 sm-k 前缀和的计数 cnt (可能为 0) 累加到最终结果中, 表示以当前元素结尾的和为 k 的子数组有 cnt 个
5. 最后将 sm 前缀和的计数加 1
另外注意, 不能先更新计数字典, 再累加结果 (即上面第 4 和第 5 步交换顺序), 因为 k 有可能是 0, 这时如果先更新的话就会错误地把每个从下标 0 开始的子数组计入到最终结果里, 显然结果就不对了…
下面代码中对上述每个步骤都有详细注释, 方便大家理解

复杂度

时间复杂度 O(N): 只需要遍历数组一遍
空间复杂度 O(N): 计数字典最多包含 N 个元素

代码

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # 前缀和+计数字典
        sm = 0
        cnts = collections.Counter()
        cnts[0] = 1
        res = 0
        for x in nums:
            # 更新前缀和
            sm += x
            # 累加sm-k前缀和的计数
            res += cnts[sm - k]
            # 更新sm前缀和的计数
            cnts[sm] += 1
        return res