974. Subarray Sums Divisible by K

题目链接

974. Subarray Sums Divisible by K

题目描述

Given an array A of integers, return the number of (contiguous, non-empty) subarrays 
that have a sum divisible by K.

 

Example 1:

Input: A = [4,5,0,-2,-3,1], K = 5
Output: 7
Explanation: There are 7 subarrays with a sum divisible by K = 5:
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]
 

Note:

1 <= A.length <= 30000
-10000 <= A[i] <= 10000
2 <= K <= 10000

题目意思就是找出所有的子集，该子集能被K整除。
一说到子集问题，我们就应该想到动态规划，而对应连续子集和问题最常用的方式就是 dp[j]表示以j为结尾的子集的和，也就是dp[0] + dp[1] +... dp[j]的和,下面的话摘自该题目的solution

As is typical with problems involving subarrays, we use prefix sums to add each subarray.
 Let P[i+1] = A[0] + A[1] + ... + A[i]. Then, each subarray can be written as P[j] - P[i] (for j > i). 

思路如下：

1. dp[j]表示从头开始到以j为结尾的子集的和，也就是dp[0] + dp[1] +... dp[j]的和
2. 根据上述思路，那么有dp[i,j] = dp[j] - dp[i]
3. 并且，对于任何一个子集都能写成如下的形式dp[i] = a * K + rem，
4. 我们假设dp[i] = a1 * K + rem1，dp[j] = a2*K + rem2，那么dp[i,j] = dp[j] - dp[i] = (a2 - a1) * K + (rem2 - rem1)
5. 如果想让dp[i,j]能被K整除，很明显，我们只需要保证rem2 - rem1 = 0即可，也就是模K的余数相同。

解决方法如下：

1. 建立mod数组，数组长度为K，mod[i]表示原数组中模K之后余数为i的数个数。
2. 遍历mod数组，如果mod[i]大于1，那么res += (mod[i] - 1) * (mod[i]) / 2,相当于里边任意的两个组合相减之后余数结果为0。
3. 最后结果加上mod[0]，表示dp[i]原本自己就能被K整除，不需要和别人相减了。

代码如下

class Solution {
    public int subarraysDivByK(int[] A, int K) {
        int sum = 0;
        int length = A.length;
        int[] mod = new int[K];
        Arrays.fill(mod, 0);
        for(int i = 0; i < length; i++) {
        	sum += A[i];
        	mod[((sum % K) + K) % K]++; // as the sum can be negative, taking modulo twice
        }
        
        sum = 0;
        for(int i = 0; i < K; i++) {
        	if(mod[i] > 1) {
        		sum += (mod[i] - 1) * mod[i] / 2;
        	}
        }
        sum += mod[0];
        return sum;
    }
}