370子区间平均值问题

题解

问题描述

给定一个长度为 n 的整数数组 arr，以及一个有理数 u/v。需要找出数组中有多少个连续的子区间，其平均值恰好等于 u/v。

解题思路

1. 转化问题：要求子区间的平均值等于 u/v，即子区间的总和等于 (u/v) * 长度。为了简化计算，我们可以将数组的每个元素都减去 u/v，然后问题转化为找出有多少个子区间的和为 0。

2. 前缀和与哈希表：

   • 使用前缀和 s 来记录从起始到当前元素的累积和。
   • 使用哈希表 d 来存储每个前缀和出现的次数。
   • 遍历数组，对于每个位置 i，计算当前的前缀和 s[i]：
     • 如果在之前的某个位置 j（j < i），s[j] 与 s[i] 相等，则子区间 (j+1, i) 的和为 0。
     • 因此，结果 res 增加 d[s[i]] 的值。
     • 然后更新 d[s[i]] 的计数。

代码解析

from itertools import accumulate
from collections import defaultdict

def solution(n: int, u: int, v: int, arr: list) -> int:
    # 将每个元素减去 u/v
    a = [x - u / v for x in arr]
    # 计算前缀和
    s = list(accumulate(a))
    # 哈希表记录前缀和出现的次数
    d = defaultdict(int)
    d[0] = 1
    res = 0
    for i in range(n):
        # 如果当前前缀和 s[i] 已存在于哈希表中，则增加相应的计数
        res += d[s[i]]
        # 更新哈希表
        d[s[i]] += 1
    return res

复杂度分析

• 时间复杂度: O(n)，其中 n 为数组的长度。
• 空间复杂度: O(n)。