Leetcode 974:和可被K整除的子数组

题目描述

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

示例1

输入：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

解题思路

连续子数组问题一般使用前缀和解决，首先要明确：假设A中有三个位置p....q.....r，其中A[p] % K = 2，A[q] % K = 2，A[r] % K = 2;则p与q之间的元素A[p+1].....A[q-1]之和一定可以被K整除，q与r之间的元素A[q+1]......A[r-1]之和一定可以被K整除，p与r之间的元素A[p+1]......A[r-1]之和一定可以被K整除。也就是三个点有个被K整除的子数组。依此类推，我们只需记录每个前缀和模K余数相等的个数，然后依此计算，需要注意的是对于恰好整除K的前缀和，除了两两组合外，单独也可以被K整除。

int subarraysDivByK(vector<int>& A, int K) {
        int sum,ans;
        int count[K] = {0};
        sum = ans = 0;
        for(int i = 0;i < A.size();i++){
            sum += A[i];
            count[(sum%K + K) % K]++;
        }
        ans += count[0];
        for(int i=0;i<K;i++){
            ans += (count[i] * (count[i] - 1)) / 2;
        }
        return ans;
    }