【算法】leetcode 974. 和可被 K 整除的子数组（前缀和思想）

前言

前缀和相关：

• leetcode 974. 和可被 K 整除的子数组（前缀和思想）
• leetcode 560. 和为K的子数组（前缀和思想II）

问题来源

• 974. 和可被 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

---

大佬解析

• 又见【前缀和】| 含 mod 运算的详细推导

---

代码

"""
需求：
    给定一个整数数组 A，返回其中元素之和可被 K 整除的（连续、非空）子数组的数目。
    
思想：
    前缀和取模
    （连续、非空）子数组 --> 相同去摸前缀和的位置组合数
    数组标记每个取模值出现频次
    组合数的意义

步骤：
    初始化一个数组，用来存储每个取模值出现频次
    循环计算前缀和取模值，更新频次表
    根据组合数意义，更新当前有效子数组数目
"""

class Solution:
    def subarraysDivByK(self, A, K: int) -> int:
        """前缀和解法"""
        count = 0  # 统计相同的前缀和数量
        presum = 0  # 前缀和初始化为0
        count_arr = [0] * K
        count_arr[0] = 1  # 提前放入 0:1，表示 求第 0 项前缀和之前，前缀和 mod K 等于 0 已经出现了 1 次(此处与presum=0的初始化操作相呼应)

        for i in A:
            presum = (presum + i) % K
            count += count_arr[presum]  # 体现组合的思想
            count_arr[presum] += 1

        return count