和可被 K 整除的子数组

标签: 数组 哈希表 前缀和

难度: Medium

给定一个整数数组 nums 和一个整数 k ,返回其中元素之和可被 k 整除的(连续、非空) 子数组 的数目。

子数组 是数组的 连续 部分。

示例 1:

输入:nums = [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]

示例 2:

输入: nums = [5], k = 9
输出: 0

提示:

  • 1 <= nums.length <= 3 * 104
  • -104 <= nums[i] <= 104
  • 2 <= k <= 104

Submission

运行时间: 48 ms

内存: 19.1 MB

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        record = {0:1}
        total = 0
        for num in nums:
            total += num
            m = total % k
            record[m] = record.get(m, 0) + 1
        ans = 0
        for k,v in record.items():
            ans += v * (v-1) // 2
        return ans

Explain

此题解采用了前缀和与哈希表的结合思想。首先,遍历数组,计算每个位置的前缀和,并对k取余,得到余数。使用哈希表record记录每个余数出现的次数。由于子数组的和可以表示为两个前缀和的差,如果两个前缀和对k取余相同,那么它们的差也能被k整除。因此,对于每个余数,从对应的次数中选取两个不同的前缀和(即组合数C(n,2)),累加到结果中。

时间复杂度: O(n)

空间复杂度: O(k)

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        record = {0:1}  # 初始化哈希表,记录余数0出现1次
        total = 0  # 前缀和
        for num in nums:
            total += num  # 计算前缀和
            m = total % k  # 对k取余
            record[m] = record.get(m, 0) + 1  # 更新哈希表
        ans = 0
        for k,v in record.items():
            ans += v * (v-1) // 2  # 计算组合数并累加到结果中
        return ans

Explore

在解法中,哈希表的键表示数组中前缀和对整数k取余后得到的余数。选择余数作为键是因为根据数学性质,如果两个数的差能够被k整除,那么这两个数对k取余的结果必定相同。因此,通过记录每个余数出现的次数,可以快速计算出有多少对前缀和的差能被k整除。

在哈希表`record`中初始化{0:1}是为了处理那些从数组起始位置直到某一点的子数组和本身就是k的倍数的情况。这里的1代表数组的和从开始到当前位置(包括整个数组)取余k为0的情况已经出现1次。这样做是为了确保可以统计包括数组首元素在内的所有符合条件的子数组。

在更新哈希表时选择对已有的值进行累加是因为我们需要记录每个余数出现的总次数。每次计算得到一个新的前缀和的余数时,我们需要增加该余数出现的计数,以便后续可以计算出从数组中任意位置到另一位置形成的子数组和是否能被k整除。直接赋值会丢失之前该余数出现的频次信息。

组合数公式`v * (v-1) // 2`用于计算从v个相同余数中选择两个(不同的前缀和位置)以形成子数组的方法数,这是基于组合数学中的组合公式C(n, 2) = n(n-1)/2。在算法中,这个公式用来计算每一种余数对应的所有可能的前缀和对的数量,这些对的差能够被k整除。因此通过该公式可以直接得到满足条件的子数组的数量。