连续子串和的整除问题 | 豆包MarsCode AI刷题

问题描述

小M是一个五年级的小学生，今天他学习了整除的知识，想通过一些练习来巩固自己的理解。他写下了一个长度为 n 的正整数序列 a_0, a_1, ..., a_{n-1}，然后想知道有多少个连续子序列的和能够被一个给定的正整数 b 整除。你能帮小M解决这个问题吗？

解题思路

1. 理解问题：

   • 我们需要找出给定序列中有多少个连续子序列的和能被一个给定的正整数 b 整除。

2. 前缀和的概念：

   • 前缀和是一个数组中从开始到当前位置的所有元素的和。例如，对于序列 [1, 2, 3]，前缀和数组为 [1, 3, 6]。
   • 通过前缀和，我们可以快速计算任意子序列的和。

3. 模运算的应用：

   • 如果两个前缀和的模 b 相同，那么这两个前缀和之间的子序列的和就能被 b 整除。
   • 例如，如果 prefix_sum[i] % b == prefix_sum[j] % b，那么 sequence[i+1] 到 sequence[j] 的子序列和能被 b 整除。

4. 哈希表的使用：

   • 我们可以使用一个哈希表（字典）来记录每个前缀和模 b 出现的次数。
   • 这样可以在遍历过程中快速计算出符合条件的子序列数量。

算法步骤

1. 初始化：

   • 初始化一个变量 prefix_sum 用于存储当前的前缀和。
   • 初始化一个哈希表 mod_count，用于记录每个前缀和模 b 出现的次数。特别地，初始化 mod_count[0] = 1，因为前缀和为 0 的情况已经存在一次。

2. 遍历序列：

   • 对于序列中的每个元素 num，更新 prefix_sum。
   • 计算当前前缀和模 b 的值 current_mod。
   • 如果 current_mod 已经在 mod_count 中，说明存在子序列和能被 b 整除，将 result 增加 mod_count[current_mod] 的值，并将 mod_count[current_mod] 加 1。
   • 如果 current_mod 不在 mod_count 中，将其初始化为 1。

3. 返回结果：

   • 最后返回 result，即符合条件的子序列数量。

总结

通过使用前缀和和模运算，我们可以高效地计算出符合条件的子序列数量。哈希表的使用使得我们可以在 O(n) 的时间复杂度内解决这个问题。 代码

def solution(n, b, sequence):
    # 初始化前缀和数组和哈希表
    prefix_sum = 0
    mod_count = {0: 1}  # 初始化模为0的情况
    result = 0

    for num in sequence:
        # 计算当前前缀和
        prefix_sum += num
        # 计算当前前缀和模 b 的值
        current_mod = prefix_sum % b
        
        # 如果当前模已经在哈希表中，说明存在子序列和能被 b 整除
        if current_mod in mod_count:
            result += mod_count[current_mod]
            mod_count[current_mod] += 1
        else:
            mod_count[current_mod] = 1

    return result

if __name__ == "__main__":
    # 你可以添加更多测试用例
    sequence = [1, 2, 3]
    print(solution(3, 3, sequence) == 3)