974. 和可被K整除的子数组

974. 和可被K整除的子数组

Scroll Down

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

思路

用了动态规划发现有一个测试样例过不去,OOM

采用了前缀和的方法,如果令 $P[i] = A[0] + A[1] + ... + A[i]$。那么每个连续子数组的和 $\textit(i, j)$ 就可以写成 $P[j] - P[i-1]$,结果等价于判断$(P[j]-P[i-1]) mod K == 0$,就等于$P[j] mod K == P[i - 1] mod K$,所以只需要记录数组和是否能够整除即可,同时记录结果。

解答

class Solution {
    // public int subarraysDivByK(int[] A, int K) {
    //     if (A == null || A.length == 0) return 0;

    //     int length = A.length;
    //     int result = 0;
    //     int[][] dp = new int[length][length];
    //     for (int i = 0; i < length; i++) {
    //         dp[i][i] = A[i];
    //         if (dp[i][i] % K == 0) result++;
    //     }
    //     for (int i = 0; i < length; i++) {
    //         for (int j = i + 1; j < length; j++) {
    //             dp[i][j] = dp[i][j - 1] + A[j];
    //             if (dp[i][j] % K == 0) result++;
    //         }
    //     }

    //     return result;
    // }

    public int subarraysDivByK(int[] A, int K) {
        Map<Integer, Integer> record = new HashMap<>();
        record.put(0, 1);
        int sum = 0, ans = 0;
        for (int num : A) {
            sum += num;
            int modulus = (sum % K + K) % K;
            int same = record.getOrDefault(modulus, 0);
            ans += same;
            record.put(modulus, same + 1);
        }
        return ans;
    }
}