Leetcode 974. 和可被 K 整除的子数组

连续+求和 用前缀

题目介绍

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

方法一：暴力破解法

• 遍历所有子数组
• 求和，找到所有能被K整除的子数组

复杂度分析

• 时间复杂度$O(n^3)$
• 空间复杂度$O(1)$

def subarraysDivByK(A, K):
    n = len(A)
    count = 0 #记录能被K整除的子数组的个数 
    for i in range(n):
        for j in range(i, n):
            if sum(A[i:j+1])%K == 0:
                count += 1
    return count

方法二：暴力破解的优化版本

暴力破解法每次对子数组求和时都要重复对数组内的元素求和，可以用前缀和来减少求和操作, 以空间复杂度来换取时间复杂度。
前缀和介绍

• 数组前0项到当前项的总和
• $preSum[0]$: 数组第0项到第0项的总和
• $preSum[1]$: 数组第0项到第1项的总和
• $preSum[2]$: 数组第0项到第2项的总和
• $$preSum[i]=A[0]+A[1]+\cdot \cdot \cdot +A[i]$$
• 数组某项，可以表示为相邻前缀和只差：$$A[i]=preSum[i]-preSum[i-1]$$
• A[i]+A[i+1]+\cdot\cdot\cdot+A[j] = preSum[j] - preSum[i-1]
  复杂度分析
• 时间复杂度$O(n^2)$
• 空间复杂度$O(n)$

def subarraysDivByK(A, K):
    preSum = [] #前缀和
    Sum = 0  #每个前缀和元素
    count = 0 #可以整除的子数组的数量
    n = len(A)
    for i in range(n):
        Sum = Sum + A[i]
        preSum.append(Sum)
    for i in range(n):
        for j in range(i, n):
            if (preSum[j] - preSum[i-1])%K == 0:
                count += 1
     return count

方法三：前缀和+哈希表

判断子数组的和能否被 K 整除就等价于判断 $(preSum[j]-preSum[i-1])\%K==0$, 根据同余定理，只要$preSum[j]\%K==preSum[i-1]\%K$,就可以保证上面的等式成立。
因此我们只需要求得所有的前缀和，并对它们取余并记录余数，每有两个前缀和余数相同，即有一个子数组满足题目要求。
首先初始化哈希表，需要注意的是哈希表初始化时要考虑前缀自身被K整除的情况，即：初始化record = {0 : 1}，然后遍历数字内的所有元素，求得前缀和preSum[i]与K的余数，哈希表中此余数出现的次数之和即为符合情况的结果。

def subarraysDivByK(A, K):
    record = {0: 1}
    total, ans = 0, 0
    for elem in A:
        total += elem
        res = total % K
        same = record.get(res, 0)
        ans += same
        record[res] = same + 1
    return ans