370.子区间平均值问题 ｜ 豆包MarsCode AI 刷题

问题描述

小U有一个长度为n的整数数组，并且选择了一个有理数 u/v。他现在想知道这个数组中有多少个连续的子区间，其平均值恰好等于 u/v。数组的子区间是指数组中一段连续的元素。

例如，给定数组 [2, 4, 1, 3, 2, 3] 和有理数 5/2，我们需要找出所有平均值等于 5/2 的子区间。

解题思路

1. 前缀和数组的计算：

   • 初始化一个前缀和数组 prefix_sum，长度为 n + 1。
   • 遍历数组 arr，计算每个位置的前缀和。

2. 目标前缀和的计算：

   • 对于每个位置 i，计算 target = prefix_sum[i] * v - i * u。

3. 哈希表的使用：

   • 使用哈希表 count_map 存储每个 target 出现的次数。
   • 在遍历过程中，检查 target 是否已经存在于 count_map 中，如果存在，则累加结果。
   • 更新 count_map 中 target 的值。

算法步骤

1. 计算前缀和：
   遍历数组 arr，计算前缀和数组 prefix_sum。
2. 遍历前缀和数组：
   对于每个 i，计算 target = prefix_sum[i] \cdot v - i \cdot u。
3. 使用哈希表：
   使用哈希表 count_map 存储每个 target 出现的次数。
4. 更新结果：
   如果 target 已经存在于 count_map 中，说明存在满足条件的子区间，将 count_map[target] 的值累加到结果中。
5. 更新哈希表：
   将当前 target 的值更新到 count_map 中。

代码实现

def solution(n: int, u: int, v: int, arr: list) -> int:
    # 初始化前缀和数组和哈希表
    prefix_sum = [0] * (n + 1)
    count_map = {}
    result = 0

    # 计算前缀和
    for i in range(1, n + 1):
        prefix_sum[i] = prefix_sum[i - 1] + arr[i - 1]

    for i in range(n + 1):
        # 计算目标前缀和
        target = prefix_sum[i] * v - i * u

        # 检查哈希表中是否存在目标前缀和
        if target in count_map:
            result += count_map[target]

        # 更新哈希表
        if target in count_map:
            count_map[target] += 1
        else:
            count_map[target] = 1

    return result

关键步骤解释

1. 前缀和计算：prefix_sum[i] 表示 arr 从 0 到 i-1 的和。
2. 目标前缀和计算：target = prefix_sum[i] * v - i * u 是为了将平均值条件转换为前缀和的形式。
3. 哈希表的使用：count_map 用于存储每个前缀和出现的次数，以便快速查找满足条件的子区间。

时间复杂度

1. 前缀和计算：

   • 计算前缀和数组 prefix_sum 需要遍历整个数组 arr 一次，时间复杂度为 O(n)。

2. 遍历前缀和数组：

   • 遍历前缀和数组 prefix_sum 需要遍历 n + 1 个元素，时间复杂度为 O(n)。
   • 在每次遍历中，计算 target 的时间复杂度为 O(1)。
   • 检查哈希表 count_map 和更新哈希表的时间复杂度为 O(1)。

因此，总的时间复杂度为 O(n)。

空间复杂度

1. 前缀和数组：

   • 前缀和数组 prefix_sum 的长度为 n + 1，空间复杂度为 O(n)。

2. 哈希表：

   • 哈希表 count_map 最多存储 n + 1 个不同的 target 值，空间复杂度为 O(n)。

因此，总的空间复杂度为 O(n)。

总结

• 时间复杂度：O(n)
• 空间复杂度：O(n)