题目描述

给定一个整数数组 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 <= A.length <= 30000
• -10000 <= A[i] <= 10000
• 2 <= K <= 10000

算法

(哈希表，前缀和) $O(n)$

1. 很容易可以发现，如果我们求出了前缀和数组 s，如果原数组区间 [j, i] 能被 K 整除，则 (s[i] - s[j - 1]) mod K == 0，即 s[j - 1] mod K == s[i] mod K。
2. 故我们只需要记录 i 之前，有多少位置的 s[j] 在模 K 后的值和 s[i] mod K 相等，这就需要用一个哈希表。
3. 具体来看，我们开一个哈希表，初始化 hash[0] = 1；然后遍历每一个位置 i，求出 s[i] mod K，如果哈希表中已经有了这个值，则总答案累计上 hash[s[i] mod K]。然后将 hash[s[i] mod K] 自增 1。

时间复杂度

• 哈希表的查询和插入均为常数，所以只需要遍历一次数组，故总时间复杂度为 $O(n)$。

空间复杂度

• 前缀和数组可以省略，只需要额外哈希表的空间，哈希表最多存放 $K$ 个数，故总空间复杂度为 $O(K)$。

C++ 代码

class Solution {
public:
    int subarraysDivByK(vector<int>& A, int K) {
        int n = A.size(), ans = 0;
        unordered_map<int, int> hash;

        hash[0] = 1;
        int s = 0;
        for (int i = 0; i < n; i++) {
            s = ((s + A[i]) % K + K) % K;

            if (hash.find(s) == hash.end())
                hash[s] = 0;

            ans += hash[s];
            hash[s]++;
        }

        return ans;
    }
};