原题传送门：《leetcode974. 和可被 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 <= A.length <= 30000
-10000 <= A[i] <= 10000
2 <= K <= 10000

AC代码：

class Solution {
public:
    int* counter;
    int subarraysDivByK(vector<int>& A, int K) {
        init(A, K);
        count(A);
        return solve(K);
    }
    int solve(int K) {
        int ans = counter[0] * (counter[0] + 1) / 2;

        for (int i = 1; i < K; ++i) {
            int tmp = counter[i];
            ans += (tmp * (tmp - 1) / 2);
        }
        delete[]counter;
        return ans;
    }
    int init(vector<int>& A, int K) {
        A[0] %= K;
        if(A[0] < 0)
        A[0] += K; //错点2
        counter = new int[K];
        memset(counter, 0, K * 4);
        for (int i = 1; i < A.size(); ++i) {
            A[i] += A[i - 1];
            A[i] %= K;
            if (A[i] < 0) A[i] += K; //错点1
        }
        return 0;
    }
    int count(vector<int>& A) {
        for (int i = 0; i < A.size(); ++i) {
            ++counter[A[i]];
        }
        return 0;
    }
};

一开始暴力TLE了没什么好说的。

需要注意

  • 数组求和时前缀和的使用
  • 负数取模的问题