统计趣味子数组的数目

标签: 数组 哈希表 前缀和

难度: Medium

给你一个下标从 0 开始的整数数组 nums ，以及整数 modulo 和整数 k 。

请你找出并统计数组中 趣味子数组 的数目。

如果 子数组 nums[l..r] 满足下述条件，则称其为 趣味子数组 ：

  • 在范围 [l, r] 内，设 cnt 为满足 nums[i] % modulo == k 的索引 i 的数量。并且 cnt % modulo == k 。

以整数形式表示并返回趣味子数组的数目。

注意：子数组是数组中的一个连续非空的元素序列。

示例 1：

输入：nums = [3,2,4], modulo = 2, k = 1
输出：3
解释：在这个示例中，趣味子数组分别是： 
子数组 nums[0..0] ，也就是 [3] 。 
- 在范围 [0, 0] 内，只存在 1 个下标 i = 0 满足 nums[i] % modulo == k 。
- 因此 cnt = 1 ，且 cnt % modulo == k 。
子数组 nums[0..1] ，也就是 [3,2] 。
- 在范围 [0, 1] 内，只存在 1 个下标 i = 0 满足 nums[i] % modulo == k 。
- 因此 cnt = 1 ，且 cnt % modulo == k 。
子数组 nums[0..2] ，也就是 [3,2,4] 。
- 在范围 [0, 2] 内，只存在 1 个下标 i = 0 满足 nums[i] % modulo == k 。
- 因此 cnt = 1 ，且 cnt % modulo == k 。
可以证明不存在其他趣味子数组。因此，答案为 3 。

示例 2：

输入：nums = [3,1,9,6], modulo = 3, k = 0
输出：2
解释：在这个示例中，趣味子数组分别是： 
子数组 nums[0..3] ，也就是 [3,1,9,6] 。
- 在范围 [0, 3] 内，只存在 3 个下标 i = 0, 2, 3 满足 nums[i] % modulo == k 。
- 因此 cnt = 3 ，且 cnt % modulo == k 。
子数组 nums[1..1] ，也就是 [1] 。
- 在范围 [1, 1] 内，不存在下标满足 nums[i] % modulo == k 。
- 因此 cnt = 0 ，且 cnt % modulo == k 。
可以证明不存在其他趣味子数组，因此答案为 2 。

提示：

  • 1 <= nums.length <= 105
  • 1 <= nums[i] <= 109
  • 1 <= modulo <= 109
  • 0 <= k < modulo

Submission

运行时间: 136 ms

内存: 29.7 MB

class Solution:
    def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
        n = len(nums)
        res = temp = 0
        mp = [0] * (n + 1)
        mp[0] = 1
        for i in range(n):
            if nums[i] % modulo == k:
                temp = (temp + 1) % modulo
            sl = (temp - k + modulo) % modulo
            if sl <= n:
                res += mp[sl]
            mp[temp] += 1
        return res

Explain

题解采用了前缀和和哈希表的思想。首先，遍历数组，计算出以每个元素为结尾的子数组中满足条件 nums[i] % modulo == k 的元素个数的累计和 temp。然后，计算出满足条件的子数组的个数。为了避免重复计算，使用哈希表 mp 来存储每个可能的累计和出现的次数。对于每个元素，我们计算出在当前累计和下，满足 cnt % modulo == k 的累计和 sl，然后将哈希表中对应 sl 的值累加到结果 res 中。最后，更新哈希表中当前累计和的计数。

时间复杂度: O(n)

空间复杂度: O(n)

class Solution:
    def countInterestingSubarrays(self, nums: List[int], modulo: int, k: int) -> int:
        n = len(nums)
        res = temp = 0
        mp = [0] * (n + 1)
        mp[0] = 1
        for i in range(n):
            if nums[i] % modulo == k:
                temp = (temp + 1) % modulo
            sl = (temp - k + modulo) % modulo
            if sl <= n:
                res += mp[sl]
            mp[temp] += 1
        return res

Explore

使用前缀和的方法可以有效减少重复计算，提高算法效率。如果直接枚举所有可能的子数组，复杂度会是 O(n^2)，这在数据量大时会非常慢。而通过使用前缀和，结合哈希表，我们可以在 O(n) 时间内完成计算，这是因为我们可以通过前缀和快速计算任意子数组的和，而哈希表则用于快速查找满足特定条件的前缀和出现的次数。

在这个算法中，哈希表用于存储前缀和的模结果及其出现的次数。哈希表的键 'key' 表示从数组开始到当前位置的子数组的和对 modulo 的余数，而哈希表的值 'value' 表示这种特定的余数出现的次数。通过记录这些信息，我们可以快速查找到前面有多少个子数组的和满足特定条件，从而快速统计满足条件的子数组个数。

这一操作的目的是更新或记录当前前缀和模 modulo 的结果的出现次数。每遍历到一个新的元素，我们都会计算新的前缀和的模结果，然后在哈希表中增加这个模结果出现次数的记录。这样，当我们后面需要查找这个前缀和模结果出现的次数时，可以直接从哈希表中得到，从而快速计算出从当前位置向前有多少子数组的和满足特定条件。

这一计算步骤是用来找出符合条件的前缀和的关键。在此算法中，我们需要找出满足 `(当前前缀和 - 某个早先前缀和) % modulo == k` 的子数组数目。通过 rearranging，我们可以得到 `当前前缀和 % modulo - k == 早先前缀和 % modulo`。因此，我们计算 `(temp - k + modulo) % modulo` 来得到应当与之前记录的前缀和模结果相等的值，这样我们就可以使用哈希表快速定位并计算满足条件的子数组的数量。