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

35 阅读2分钟

问题文本

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

问题分析

该问题的暴力解法是对每一个连续子序列作判断,一个长度为n的序列共有n(n-1)/2个连续子序列,每一个子序列的和的计算复杂度为子序列长度,复杂度为O(n^3)。

我们发现在暴力解法中,一个子序列的和其实可以由其余子序列求出。我们采用前缀和的方式简化子序列求和,使用pre数组存储前缀和,pre[i]表示前i个数字的和。有以下公式:

i=j+1z=pre[z]pre[j],0<=j<z<=n\sum_{i=j+1}^{z} =pre[z]-pre[j], \quad 0 <= j < z <= n

通过前缀和方式,我们可以将求和复杂度降低到O(1)。我们可以从前往后遍历数组并同步更新前缀和数组,对每一个被遍历的数字,对所有以该数字结尾的连续子序列作判断,每次判断的复杂度为该数字位置,此时时间复杂度为O(n^2),空间复杂度为O(n),代码如下:

def solution(n, b, sequence):
    if n == 0:
        return -1
    preSum = [0] * (n + 1)
    for i in range(1, n + 1):
        preSum[i] = preSum[i - 1] + sequence[i - 1]
    ret = 0
    for i in range(n + 1):
        for j in range(i+1, n + 1):
            if not (preSum[j] - preSum[i]) % b:
                ret += 1
    return ret

我们重申问题是要求有多少个连续子序列的和能够被一个给定的正整数 b 整除。而我们知道,对于同一个除数,相同模数的两个数字相减的结果可以被除数整除。

因此我们可以进一步优化上述方案,首先借助前缀和实现快速求和,之后使用字典modCount存储前缀和对于除数b的模数的出现次数。当遇到一个新的数字时,我们先求出当前前缀和对除数b的模数,此时模数与当前模数相同的前缀和数量即为可以被正整数b整除的子序列个数。

由于只需要遍历一遍数组,时间复杂度为O(n),对于前缀和我们只需要保存当前前缀和即可,空间复杂度为O(1),代码如下:

def solution(n, b, sequence):
    if n == 0:
        return -1
    modCount = {0:1}
    preSum = 0
    ret = 0
    for num in sequence:
        preSum += num
        mod_value = preSum % b

        if mod_value in modCount:
            ret += modCount[mod_value]
            modCount[mod_value] += 1
        else:
            modCount[mod_value] = 1
    return ret