esc to dismiss

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]

解体思路

使用 前缀和 + 哈希表 的算法。

我写的代码没有官方解法的简洁。官方解法:

func subarraysDivByK(A []int, K int) int {
    record := map[int]int{0: 1}
    sum, ans := 0, 0
    for _, elem := range A {
        sum += elem
        modulus := (sum % K + K) % K
        ans += record[modulus]
        record[modulus]++
    } 
    return ans
}

思路是前缀和,对于前缀和 pre[[i] = A[0] + A[1] ... + A[i], 那么计算子数组的和 sum(i, j) 即可以转换为: pre[j] - pre[i-1],那么子数组和能否被 K 整除即变成 ( pre[j] - pre[i-1]) % K 0, 根据 [同余定理]] ,这个等式可以转换为 pre[j] 和 pre[i-1] 具有相同的余数,即 pre[j] % K pre[i-1] % K

这样,本题就可以转换为 前缀和+哈希表的应用,与 560. 和为K的子数组 题解法就类似了。

要注意的是:元素中有负数,那么取余也会得到负数,需要转换成正数余: ( sum % K + K) % K, 否则只需要 sum % K

经验教训

本题要考虑的点较多,通过 LeetCode 提交代码失败了几次,都是因为用例没过,究其原因是场景考虑不全,这在考试中是大忌。

我在 LeetCode 上提交的答案:

func subarraysDivByK(A []int, K int) int {
    size := len(A)
    count := 0
    pre := 0
    m := make(map[int]int)
    remain := 0
    m[0] = 1

    for i := 0; i < size; i++ {
        pre += A[i]
        if K != 0 {
            remain = pre%K
        }

        count += m[remain]
        if remain < 0 {
            count += m[K+remain]
            m[K+remain] += 1
        } else if remain > 0 {
            count += m[-remain]
            m[remain] += 1
        } else {
            m[remain] += 1
        }
    }

    return count
}

主要是没有考虑到余数为负的情况。

x