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

前缀和+HashMap

常规思路是：首先求得前缀和数组，然后检查每一个子数组的和是否能被K整除。这样做的时间复杂度是O(n^2)，会超时，需要继续优化。优化的思路是不用检查每个子数组，而是能快速找到和能被K整除的数组。可以使用HashMap，前缀和sum[i]除以K的余数为key，该余数出现的次数为value，如果有两个余数相同的前缀和，那么它们对应的区间和就能被K整除。遍历数组，依次将对应的entry加入HashMap，记当前位置的前缀和除以K的余数为r，则前面有多少个对应前缀和余数同为r的位置，则符合条件的子数组就新增几个。

需要注意的是两个前缀和相减对应的区间不能从下标0开始，如果当前位置的sum[i]刚好就能被K整除，则符合条件的子数组个数加1。

class Solution {
    public int subarraysDivByK(int[] A, int K) {
        int count=0,n=A.length,s=0;
        int[] sum=new int[n];
        for(int i=0;i<n;++i){
            s+=A[i];
            sum[i]=s;
        }
        int r=0;//r为余数
        //key为sum[i]除以K的余数，value为该余数出现的次数
        HashMap<Integer,Integer> map=new HashMap<>();
        for(int i=0;i<n;++i){
            r=(sum[i]%K+K)%K;//防止余数为负数
            if(r==0) count+=1;//如果sum[i]可以被K整除，count加1

            if(map.containsKey(r)){
                count+=map.get(r);
                map.put(r,map.get(r)+1);
            }
            else map.put(r,1);
        }
        return count;
    }
}

原题链接：https://leetcode-cn.com/problems/subarray-sums-divisible-by-k