前缀和及简单应用

1、思想　　

　　通常，涉及连续子数组问题的时候，我们使用前缀和来解决。

　　我们令$P[i]=A[0]+A[1]+\ldots+A[i]$，那么每个连续子数组的和$\operatorname{sum}(i, j)$ 就可以写成 $P[j]-P[i-1]$（其中 0 < i < j）的形式。

2、应用

a、和为k的子数组

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

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

输出: 2 , [1,1] 与 [1,1] 为两种不同的情况。

思考：

「[j..i] 这个子数組和为 $k$ 」这个条件我们可以转化为$\operatorname{pre}[i]-\operatorname{pre}[j-1]==k$

所以我们考虑以 i 结尾的和为 k 的连续子数组个数时只要统计有多少个前缀和为 $\operatorname{pre}[i]-k$ 的 $\operatorname{pre}[j]$ 即可。我们建立哈希表 $\operatorname{maps}$，以和为键，出现次数为对应的值，记录 $\operatorname{pre}[i]$ 出现的次数，从左往右边更新 $\operatorname{maps}$ 边计算答案，那么以 i 结尾的答案 $\operatorname{maps}[\operatorname{pre}[i]-k]$ 即可在 O(1) 时间内得到。最后的答案即为所有下标结尾的和为 k 的子数组个数之和。

代码

class Solution:
    # 暴力
    # def subarraySum(self, nums, k):
    #     res = 0
    #     if not nums:
    #         return res

    #     for i in range(0, len(nums)):
    #         s = 0
    #         path = []
    #         for j in range(i, len(nums)):
    #             s += nums[j]
    #             path.append(nums[j])
    #             if s == k:
    #                 res += 1

    #     return res

    def subarraySum(self, nums, k):
        res = 0
        if len(nums) == 0:
            return res
        
        maps = {0:1}
        pre_num = 0

        for i in nums:
            pre_num += i
            if maps.get(pre_num-k):
                res += maps[pre_num-k]
            if maps.get(pre_num):
                maps[pre_num]+=1
            else:
                maps[pre_num] = 1
            
        return res

View Code

b、和可被k整除的数组

给定一个整数数组 A，返回其中元素之和可被 K 整除的（连续、非空）子数组的数目。

示例：

输入：A = [4,5,0,-2,-3,1], K = 5
输出：7
解释：
有 7 个子数组满足其元素之和可被 K = 5 整除：
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

思考：

判断子数组的和能否被 $K$ 點除就等价于判断 $(P[j]-P[i-1]) \bmod K==0$, 根据同余定理, 只要 $P[j] \bmod K==P[i-1] \bmod K$,就可以保证上面的等式成立。

　　因此我们可以考虑对数组进行遍历，在遍历同时统计答案。当我们遍历到第 i 个元素时，我们统计以 i 结尾的符合条件的子数组个数。我们可以维护一个以前缀和模 K 的值为键，出现次数为值的哈希表 $\operatorname{record}$，在遍历的同时进行更新。这样在计算以 i 结尾的符合条件的子数组个数时，根据上面的分析，答案即为 [0..i−1] 中前缀和模 K 也为 $P[i] \bmod K$ 的位置个数，即 $\operatorname{record}[P[i] \bmod K]$。

　　最后的答案即为以每一个位置为数尾的符合条件的子数组个数之和。需要注意的一个边界条件是，我们需要对哈希表初始化，记录 $\operatorname{record}[0] = 1$，这样就考虑了前缀和本身被 K 整除的情况。

代码：

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        record = {0: 1}
        total, ans = 0, 0
        for elem in A:
            total += elem
            modulus = total % K
            same = record.get(modulus, 0)
            ans += same
            record[modulus] = same + 1
        return ans

View Code