Leetcode 974 Subarray Sums Divisible by 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]

提示：

  1. 1 <= A.length <= 30000
  2. -10000 <= A[i] <= 10000
  3. 2 <= K <= 10000

思路分析：
当我们看到连续子数组问题的时候，我认为我们必须马上要想到前缀数组。这里由于是考虑连续子数组的和，那么我们就用前缀和数组来分析。

不妨设pre[i]表示sum(A[:i])，那么pre[0] = 0, pre[1] = A[0]
现在我们要考虑某段子数组A[i:j]的和是否能整除K，只需要判断(pre[j]-pre[i]) % K == 0。但是如果暴力法判断所有的i,j，时间复杂度会达到O(n^2)，肯定是不能暴力法的。
我们注意到想让(pre[j]-pre[i]) % K == 0只需要让pre[j] % K == pre[i] % K即可。我们使用一个dp[i]表示目前模K余i的数的个数，默认dp[0] = 1

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def subarraysDivByK(self, A, K):
        n = len(A)
        pre = [0] * (n+1)
        for i in range(1, n+1):
            pre[i] = pre[i-1] + A[i-1]

        dp = [0] * K
        dp[0] = 1
        res = 0
        for i in range(1, n+1):
            val = pre[i] % K
            res += dp[val]
            dp[val] += 1
        return res

# 73 / 73 test cases passed.
# difficultu: medium
# Runtime: 288 ms