计数具有与子数组长度相同的模K的子数组

简介

在一个给定的整数数组中，我们需要统计所有长度为K的子数组（连续的K个元素），并且满足子数组元素之和模K的余数等于子数组长度K的模K余数的个数。

问题分析

我们可以首先预处理一下前缀和数组，即 prefix_sum[i] 表示原数组从0到i的元素和。那么题目要求的模K余数可以改写为 (prefix_sum[j] - prefix_sum[i-1]) % K，即每个长度为K的子数组的模K余数都可以转化为两个前缀和的模K余数之差。

解法一

我们用一个哈希表 count 来存储每个模K余数出现的次数。遍历前缀和数组，每遍历到一个新的前缀和值，就把它对K取模得到模K余数，然后在哈希表中查找这个模K余数出现的次数，并把次数加入到答案中。同时我们再把这个模K余数的出现次数加1，表示我们又增加了一个以这个模K余数结尾的长度为K的子数组。

具体实现如下：

def count_subarrays(nums, K):
    prefix_sum = [0]
    for num in nums:
        prefix_sum.append(prefix_sum[-1] + num)
    count = {0: 1}  # 初始化为{0: 1}，表示前缀和为0的模K余数已经出现了1次
    ans = 0
    for i in range(1, len(prefix_sum)):
        mod = prefix_sum[i] % K
        ans += count.get(mod - i % K, 0)
        count[mod - i % K] = count.get((mod - i % K), 0) + 1
    return ans

解法二

我们可以不用哈希表，而是用一个长度为K的数组 count，其中 count[j] 表示当前已经找到的子数组中，长度为K且模K余数为j的子数组个数。遍历前缀和数组，每遍历到一个新的前缀和值，就把它对K取模得到模K余数，然后在 count 数组中查找这个模K余数出现的次数，并把次数加入到答案中。同时我们再把以这个模K余数结尾的长度为K的子数组个数加1，即 count[mod] += 1。

具体实现如下：

def count_subarrays(nums, K):
    prefix_sum = [0]
    for num in nums:
        prefix_sum.append(prefix_sum[-1] + num)
    count = [0] * K  # 初始化为长度为K的0数组
    count[0] = 1     # 初始化为[1] + [0]*(K-1)，长度为K的模K余数为0的子数组已经找到了1个
    ans = 0
    for i in range(1, len(prefix_sum)):
        mod = prefix_sum[i] % K
        ans += count[mod - i % K]
        count[mod] += 1
    return ans

总结

这道题其实是一个非常典型的前缀和问题，但涉及到了一些取模的技巧。通过用哈希表或数组来记录已经找到的子数组个数，我们可以很快地得到答案。两种解法的时间复杂度都是O(N)，空间复杂度也都是O(K)。

以上是本文的介绍，完整可运行代码请见下方。