首页手记 LeetCode 974. 和可被 K 整除的子数组...

LeetCode 974. 和可被 K 整除的子数组 | Python

标签：

Python 算法

974. 和可被 K 整除的子数组

题目来源：力扣（LeetCode）https://leetcode-cn.com/problems/subarray-sums-divisible-by-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]

提示：

1 <= A.length <= 30000
-10000 <= A[i] <= 10000
2 <= K <= 10000

解题思路

思路：前缀和

首先这里要提前说明一下，关于除法取整，Python 采用的是向下取整。这里可能跟预想出现偏离的情况，比如，除数为负数的情况。看以下示例：

>>> 10 // 3
3
>>> 10 % 3
1
>>> 10 // -3
-4
>>> 10 % -3
-2

这里可以看到，关于正数取整，取模跟预想都不会有太大的偏离。但是对于负数，出现的结果可能就跟预想的不太一样。这里是因为前面说，Python 采用的是向下取整。

对于 10 ÷ -3 = -3.3333，这时向下取整就会得到 -4，那么余数就是 -2。这就是为什么会出现上面结果的原因。这部分内容仅做一些提示。

在本题当中 K，这里是除数，题目中说明【2 <= K <= 10000】，所以不会有上面所述的情况。但是 Python 当中负数取余所得的结果是正数，这里跟一些其他编程语言不同。有些语言负数取余的结果是负数，所以要额外进行处理。

本篇幅使用的是前缀和的思想。关于前缀和，大致说明下。

如果要求前 i 项的和，那么：

P[i] = A[0]+A[1]+...+A[i]

相应的前 i-1 项的和为：

P[i-1] = A[0]+A[1]+...+A[i-1]

那么，A[i] 的值也可以表示为

A[i] = P[i] - P[i-1]

那么相应的，如果要计算 i 项到 j 项连续子数组的和，也可用写成如下的形式：

sum[i...j] = P[j] - P[i-1]，其中 (0 < i < j)

那么题目中要求，判断子数组的和是否能够被 K 整除，现在就等同于判断 (P[j]-P[i-1]) mod K == 0 是否成立。

在这里，额外提及一个定理：同余定理。

同余定理：给定一个正整数 m，如果两个整数 a 和 b 满足 a-b 能够被 m 整除，即 (a-b)/m 得到一个整数，那么就称整数 a 与 b 对模 m 同余。

那么上面需要判断的式子也就可以转换为求 P[j] mod K == P[i-1] mod K 式子是否成立。

具体的方法：

维护哈希表，其中哈希表键为前缀和模 K 的值，值为出现的次数。
遍历数组每项，求当前前缀和模 K，存入哈希表中
- 当不存在表中，则将键跟值存入
- 存在时，对应键的值 +1
遍历的同时，进行统计，如果哈希表中存在 key 与当前前缀和模 k 的值相等时：
- 说明前面存在前缀和模 K 的值与此次计算的值相同
- 将满足条件的 key 出现的次数，累加到结果中

具体的代码实现如下。

代码实现

class Solution:
    def subarraysDivByK(self, A: List[int], K: int) -> int:
        # 这里是考虑前缀和本身被 K 整除的情况
        hashmap = {0: 1}
        pre_sum = 0
        cnt = 0

        for i in range(len(A)):
            pre_sum += A[i]
            # 取模
            mod = pre_sum % K
            # 这里使用字典的 get 方法
            # 当存在相同的键时，累加到 cnt
            if mod in hashmap:
                cnt += hashmap[mod]
                hashmap[mod] += 1
            # 如果键在哈希表中，则次数加 1，
            # 否则初始化为 1
            else:
                hashmap[mod] = 1
        
        return cnt