LeetCode 974.和可被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]

题目链接

题目分析

前缀和双层遍历

最直观的思路是先求出前缀和数组，然后遍历前缀和数组：

int subarraysDivByK(int* A, int ASize, int K){
    int* ans = (int*)calloc(ASize, sizeof(int));
    int count = 0;
    int temp = 0;
    for (int i = 0; i < ASize; i++) {
        temp += A[i];
        ans[count++] = temp;
    }
    int res = 0;
    for (int i = 0; i < ASize - 1; i++) {
        if (ans[i] % K == 0) res++;
        for (int j = i + 1; j < ASize; j++) {
            if ((ans[j] - ans[i]) % K == 0) res++;
        }
    }
    if (ans[ASize - 1] % K == 0) res++;
    return res;
}

这样的时间复杂度是O(n^2)，会超时。

前缀和+哈希表（同余定理）

同余定理：给定一个正整数m，如果两个整数a和b满足a-b能够被m整除，即(a-b)/m得到一个整数，那么就称整数a与b对模m同余，记作a≡b(mod m)。对模m同余是整数的一个等价关系。

也就是说，如果两个数对K的余数相同，则两数之差能被K整除。

因此，我们可以用哈希表储存数组中每个数字对K的余数，遇到相同余数则结果+1。

class Solution {
public:
    int subarraysDivByK(vector& A, int K) {
        unordered_map m = {{0, 1}};
        int sum = 0, ans = 0;
        for (auto a : A) {
            sum += a;
            int mod = (sum % K + K) % K;
            if (m.count(mod)) {
                ans += m[mod];
            }
            m[mod]++;
        }
        return ans;
    }
};