974. Subarray Sums Divisible by K**-CFANZ编程社区

974. Subarray Sums Divisible by K**

https://leetcode.com/problems/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

解题思路

Prefix Sum. 使用常规思路容易 TLE (超时). 要解决这道题, 需要了解求余的性质. 具体参考 (Python) Concise Explanation and Proof.

求余有如下性质:

(a + b) mod n = [(a mod n) + (b mod n)] mod n
(a mod n) mod n = a mod n

另外考虑到负数的情况, 如果 a < 0, 可以用 a % K + K 来处理 a 为负数的情形. 因为:

(b + a) % K = (b % K + a % K) % K
# 而
(b + a % K + K) % K = (b % K + a % K % K + K % K) % K # 性质 1
                    = (b % K + a % K + 0) % K  # 性质 2
                    = (b % K + a % K) % K
                    = (b + a) % K

在 (Python) Concise Explanation and Proof 中, 作者的想法是, 要判断 sum[i, ... j) = a[i] + a[i + 1] + .... + a[j - 1] 是否能被 K 整除, 方法是判断 sum[j] % K 是否等于 sum[i] % K. 原因在于如果 sum[i, ... j) % K == 0, 那么由于 sum[i, j) = sum[0, j) - sum[0, i), 则有:

( sum[0, j) - sum[0, i) ) == n * K # 其中 n 未知
sum[0, i) = sum[0, j) - n * K
sum[0, i) % K = ( sum[0, j) - n * K ) % K
              = ( sum[0, j) % K - n * K % K) % K
              = ( sum[0, j) % K - 0) % K
              = sum[0, j) % K

因此, 判断 sum[i, j) 是否能被 K 整除的方法是判断 sum[0, j) % K 是否和 sum[0, i) % K 相等. 具体来说, 我们可以用 unordered_map<int, int> record 来存储已经访问过的元素的 key = sum[0, i) % K, value 保存 key 的个数. 当访问到 sum[0, j) 时, 首先求 key' = sum[0, j) % K, 看 key' 能否在 record 中访问到.

还需要注意的是对负数的处理, 在求 sum[0, i) 的时候, 如果当前访问的元素 a < 0, 可以用 a % K + K 来替代. 假设 a = -3, K = 5, 在 C++ 中 a % K = -3, 而 a % K + K 的结果为 2. 但是 (b + a) % K 和 (b + a % K + K) % K 结果相等, 上面已经证明过了.

C++ 实现 1

在 How do you come up with Prefix Sum for this problem? 中的回答中, 有人说道他们能想到 Prefix Sum 方法的原因是:

Whenever I see the word Subarray immediately one of the first few thoughts that comes to my mind is prefix sum or sliding window.

果然… 得总结解题技巧.

class Solution {
public:
    int subarraysDivByK(vector<int>& A, int K) {
        // record 初始化为 {{0, 1}}, 了解代码逻辑之后,
        // 考虑如果数组 A 中只有一个元素的情况, 就知道为啥这样初始化了.
        unordered_map<int, int> record{{0, 1}};
        int count = 0, sum = 0;
        for (int i = 0; i < A.size(); ++ i) {
            sum += A[i] % K + K;
            int key = sum % K; // 只需要在 record 中记录 key 而不是 sum
            count += record[key]++;
        }
        return count;
    }
};