题目描述

[EN | CN]

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

示例:

1
2
3
4
5
输入: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 <= A.length <= 30000
  • -10000 <= A[i] <= 10000
  • 2 <= K <= 10000

解法 1:暴力窗口遍历

我们很容易想到用一个固定长度的窗口在数组上滑动,每次计算窗口内的和是否可以被 K 整除,然后窗口的长度从 1 开始到数组的长度。

这个方法容易想到,但是比较麻烦且很费时,非常容易超时。

复杂度分析略。

实现与结果如下:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:

        A = [i % K for i in A]

        ret = 0
        last_sum = 0
        for win_len in range(1, len(A) + 1):
            win_sum = last_sum + A[win_len - 1]
            last_sum = win_sum
            if win_sum % K == 0:
                ret += 1
            for i in range(len(A) - win_len):
                win_sum = win_sum - A[i] + A[i + win_len]
                if win_sum % K == 0:
                    ret += 1

        return ret
  • 执行结果:超出时间限制。

解法 2:前缀和

很酷的方法。

所谓 前缀和 指的是前 $n$ 个数的和。对于长度为 $n$ 的数组 $A$ 来说,其前缀和数组 $P$ 的长度也为 $n$,其中第 $i$ 个元素 $P_i = \sum_{j=1}^iA_i$。前缀和有很好的性质,对于任意 $i>0$、$j>0$,若 $P_i\mod{K}==P_j\mod{K}$,那么从 $i$ 到 $j$ 的数组和能够被 $K$ 整除

具体到本题,我们只需要算出前缀和数组,然后对每一个可能的余数算一下长度为 2 的组合的数量,将这些数量求和即为题目所求答案。

注意这里有个情况容易漏掉,就是当 $P_i$ 自身也有可能能被 $K$ 整除,因此我们应该 在前缀和数组的开头加个 0,从而囊括了这种情况。

实现与结果如下:

1
2
3
4
5
6
7
8
from collections import Counter
class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        P = [0]
        for x in A:
            P.append((P[-1] + x) % K)
        count = Counter(P)
        return sum([v * (v - 1) // 2 for v in count.values()])
  • 执行用时:352 ms, 在所有 Python3 提交中击败了 91.06% 的用户。
  • 内存消耗:17.5 MB, 在所有 Python3 提交中击败了 100.00% 的用户。