题目描述

给定一个数组arr和整数k，现在我们需要算出所有子数组的和，并求和后的结果需要对k取余。求最终的结果。

函数签名

def subarray_sum_mod_k(arr: List[int], k: int) -> int:
    pass

输入描述

• arr: 整数数组，长度为 n (1 <= n <= $10^4$)
• k: 整数，常数 (1 <= k <= $10^9$)

输出描述

最终求出所有子数组求和结果对k取余的结果。

例子

assert subarray_sum_mod_k([1, 2, 3], 2) == 0
assert subarray_sum_mod_k([1, 2, 3], 3) == 2
assert subarray_sum_mod_k([1, 2, 3], 4) == 1

解题思路

我们使用前缀和的思想，假设现在有一个数组 arr，现在我们想求的是以 i 为右端点的所有子数组的和，我们可以通过预处理前缀和的方法来减少时间复杂度：

$$ s_i = \sum_{j=0}^{i-1} a_j $$

现在我们来看一下如何计算左右端点分别为 i 和 j 的子数组的和，可以使用前缀和来简化计算过程：

$$ \begin{aligned} S_{i,j} &= \sum_{k=i}^j a_k = s_j - s_i \ S_{0,i} &= \sum_{k=0}^i a_k = s_i \end{aligned} $$

现在回到我们的题目中，我们要求所有子数组的和对 k 取模后的结果。我们可以先把所有子数组的和求出来，然后对 k 取模：

$$ \begin{aligned} res &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}(\sum_{k=i}^{j-1}a_k) % k \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}(\sum_{k=i}^{j-1}a_k % k) \end{aligned} $$

接下来我们需要快速求出子数组的和对 k 取模后的结果，我们可以把求和运算分成分子段，分别对每个子段求和并取模：

$$ \begin{aligned} res &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}(\sum_{k=i}^{j-1}a_k % k) \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}((\sum_{k=0}^{j-1}a_k - \sum_{k=0}^{i-1} a_k)%k) \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}(((s_j-s_i) % k+k) % k) \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}((s_j-s_i+s_i) % k) \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}((s_j%k-s_i%k+k) % k) \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}(s_j%k-s_i%k+k%k) % k \ &= \sum_{i=0}^{n-1}\sum_{j=i+1}^{n}(s_j%k-s_i%k) % k \ \end{aligned} $$

其中 $k%k = 0$。现在对于每个左右端点的区间，我们只需要计算出两个前缀和的差，再对 k 取模即可。

我们可以使用哈希表来记录前缀和余数相同的值出现的次数，然后我们在计算出所有的前缀和余数后，算出当前前缀和的余数为 x，那么当前答案就应该加上前缀和余数为 x 的数量，然后我们将当前前缀和的余数加入哈希表中。

代码实现

from collections import defaultdict

def subarray_sum_mod_k(arr: List[int], k: int) -> int:
    prefix_sum = defaultdict(int)
    prefix_sum[0] = 1 # 初始化，因为 prefix_sum[-1]=0
    res, sum_i = 0, 0
    for i in range(len(arr)):
        sum_i = (sum_i + arr[i]) % k
        if sum_i in prefix_sum:
            res += prefix_sum[sum_i]
        prefix_sum[sum_i] += 1
    return res

时间复杂度：$O(n)$

空间复杂度：$O(k)$