查找每个元素可被K整除的子序列数

介绍

在一个给定的数组中，计算有多少个子序列的元素之和可以被K整除。

举例：

给定数组 [4, 5, 0, -2, -3, 1]，整数K为5，可以得到以下子序列：

• [4]
• [4, 5, 0, -2, -3, 1]
• [5, 0]
• [5, 0, -2, -3]
• [0]
• [0, -2, -3]
• [-2, -3]
• [5]
• [5, 0, -2]
• [0, -2]
• [5, 0, -2, -3, 1]
• [0, -2, -3, 1]

其中，有5个子序列的元素之和可以被5整除，因此返回值为5。

解法

针对这种问题，经典的解法是通过前缀和来解决。

定义一个数组 prefix 用来存储前缀和，则 prefix[i] 表示数组中前i个元素的和。若想确定第i个元素到第j个元素的和，只需要计算 prefix[j] - prefix[i-1] 即可。

有了前缀和，就可以在O(n^2)的时间复杂度内计算每个子序列的元素之和。但这样会超时，需要通过优化来进行更快的计算。

观察到只有两个元素之和可以被K整除的情况，可以拆分为以下两种情况：

• 一组元素和可以被K整除，另一组元素和也可以被K整除。
• 一组元素和不可以被K整除，另一组元素和可以被K整除，那么这两组元素之和的和也可以被K整除。

例如，若K为5，存在一个子序列[4, 1]，它们的和为5，说明[4]与[1]都能被5整除，所以[4, 1]也能被5整除；而存在一个子序列[4, 5, 0, -2, -3, 1]，其中[5, 0]的和为5，说明[5, 0]不可被5整除，但[5, 0, -2, -3]的和可以被5整除，所以[5, 0, -2, -3]与[-2, -3]的和也可以被5整除。

因此，只需要计算某个元素之前的前缀和 mod K 的值，再将这些值存到一个哈希表中。对于每个元素的前缀和 mod K 的值，若哈希表中已存在相同的值，则说明当前元素所在位置可以作为子序列的末尾，这些元素之间的子序列中间的元素之和可以被K整除；否则将当前元素的前缀和 mod K 的值插入哈希表中，继续遍历数组。最终，哈希表中统计的键-值对就是满足要求的子序列的个数。

代码片段

def subarraysDivByK(A, K):
    d = {0: 1}
    s = res = 0
    for num in A:
        s = (s + num) % K
        if s in d:
            res += d[s]
        d[s] = d.get(s, 0) + 1
    return res

此代码为Python 3的实现，通过哈希表统计子序列个数。