[LeetCode] 974. 和可被 K 整除的子数组 (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

解题思路

令数组的前 $n$ 项和为 $S[n]=A[1]+\dots+A[n]$,则连续子数组 $A[i]$ 到 $A[j]$ 的和为 $S[j]-S[i-1]$,则该数组和可以被 $K$ 整除相当于 $(S[j]-S[i-1])\mod K=0$,根据同余定理,$S[j]\mod K=S[i-1]\mod K$ 成立则前式成立。因此,我们维护一个前 $n$ 项和模 $K$ 的值为键,个数为值的哈希表。

代码

Python 3.6

1
2
3
4
5
6
7
8
9
10
11
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

执行用时:376 ms

内存消耗:17.6 MB

C++

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
class Solution {
public:
int subarraysDivByK(vector<int>& A, int K) {
unordered_map<int, int> record = {{0, 1}};
int total = 0, ans = 0;
for (int elem : A) {
total += elem;
int modules = (total % K + K) % K; // C++ 取模的特殊性。当被除数为负数时取模结果为负数,需要纠正
int same = (record.count(modules)) ? record[modules] : 0;
ans += same;
record[modules] = same + 1;
}
return ans;
}
};

关于负数取模请参见 Python 与 C++ 关于负数取模的不同

执行用时:148 ms

内存消耗:30.6 MB

参考

https://leetcode-cn.com/problems/subarray-sums-divisible-by-k/solution/he-ke-bei-k-zheng-chu-de-zi-shu-zu-by-leetcode-sol/