[python]连续子串和的整除问题(前缀和) | 豆包MarsCode AI刷题

问题描述

小M正在学习整除的知识。他写下了一个长度为 nn 的正整数序列，并希望知道有多少个连续子序列的和能够被一个给定的正整数 bb 整除。

输入与输出

输入:

1. n：数组的长度。
2. b：给定的正整数。
3. sequence：长度为 n 的正整数数组。

输出: 返回可以被 b 整除的连续子序列的个数。

测试样例

样例 1

输入：

n = 3, b = 3, sequence = [1, 2, 3]

输出：

3

样例 2

输入：

n = 4, b = 5, sequence = [5, 10, 15, 20]

输出：

10

样例 3

输入：

n = 5, b = 2, sequence = [1, 2, 3, 4, 5]

输出：

6

解题思路

1. 基本思路

这个问题可以用暴力枚举的方式解决，但时间复杂度为 O(n^2)，在 n 较大时效率极低。我们需要一个更高效的算法来解决这个问题。

关键在于：

• 我们需要快速计算每个连续子序列的和。
• 并判断这个和是否能被bb 整除。

这里，我们引入前缀和与哈希表来优化。

2. 核心公式推导

连续子序列的和

假设数组$sequence=[a_0,a_1,...,a_{n−1}]$，其前缀和数组定义为：

对于数组中任意子序列 a[i:j]a[i:j]，其和可以表示为：

整除条件

如果 sum(a[i:j])%b=0 ，可以重写为：

prefix[j] % b = prefix[i−1] % b

即，当两个前缀和对 b 取模的结果相同时，以这两个前缀和结尾的子序列和能被 b 整除。

3. 代码实现

以下是基于 Python 的实现：

def solution(n, b, sequence):
    # 哈希表记录前缀和模 b 的出现次数
    mod_count = {0: 1}  # 初始值，表示前缀和本身为 0 时的情况
    prefix_sum = 0
    count = 0

    for num in sequence:
        prefix_sum += num
        mod = prefix_sum % b
        # 如果模值为负数，调整为非负数
        mod = (mod + b) % b

        # 如果当前模值已出现，增加对应的子序列个数
        count += mod_count.get(mod, 0)

        # 更新当前模值出现次数
        mod_count[mod] = mod_count.get(mod, 0) + 1

    return count

4. 示例详解

示例 1

输入：

n = 3, b = 3, sequence = [1, 2, 3]

步骤：

• 初始化 mod_count = {0: 1}。

• 遍历数组，计算前缀和模 bb：

  • 索引 0：前缀和为 1，模值为 1，mod_count = {0: 1, 1: 1}。
  • 索引 1：前缀和为 3，模值为 0，mod_count = {0: 2, 1: 1}，累加 1。
  • 索引 2：前缀和为 6，模值为 0，mod_count = {0: 3, 1: 1}，累加 2。

输出： 3（对应子序列：[3], [1, 2], [1, 2, 3]）

5. 复杂度分析

时间复杂度

• 遍历数组：O(n)
• 哈希表操作（插入与查询）：均摊 O(1)
• 总复杂度：O(n)

空间复杂度

• 哈希表存储模值的次数，最多为 b，空间复杂度为 O(b)。