求长度≥K的子数组最大平均值的算法优化求助

问题描述 #

给定一个包含n个元素的数组，定义T(i, L) = (a[i] + a[i+1] + ... + a[i+L-1])/L（其中1 ≤ i ≤ n-L），要求子段长度L ≥ K，找出T(i, L)的最大值。

输入输出要求 #

输入 #

• 第一行是两个整数N, K（1 ≤ N ≤ 3e5）
• 第二行是N个非负整数（0 ≤ a[i] ≤ 1e6）

输出 #

输出满足L ≥ K的T(i, L)的最大值，需保留足够精度（如示例保留6位小数）

示例 #

示例1 #

输入：

4 1
1 0 4 3

输出：4

示例2 #

输入：

5 2
2 4 3 4 1

输出：3.666667

现有问题 #

我目前只能实现时间复杂度为O(n²)的解法，对于n=3e5的数据规模完全无法通过，寻求更高效的算法方案。

高效解法：二分答案法 #

核心思路 #

数组元素均为非负数，子段平均值具有单调性：若存在长度≥K的子段平均值≥mid，则所有小于mid的平均值都满足条件；反之则不满足。利用这个特性，通过二分法逼近最大平均值。

具体步骤 #

1. 确定二分范围：左边界left=0，右边界right设为数组中的最大值（非负数组的子段平均值不可能超过单个元素的最大值）。
2. 可行性检查：对当前中间值mid，判断是否存在长度≥K的子段平均值≥mid：
   • 将数组每个元素减去mid，得到新数组b[i] = a[i] - mid，问题转化为是否存在长度≥K的子段和≥0。
   • 计算前缀和数组s，其中s[0]=0，s[i] = b[1]+b[2]+...+b[i]（数组从1开始索引）。
   • 遍历过程中维护j ≤ i-K的最小前缀和s[j]，若s[i] - min_s ≥0，说明存在符合条件的子段，返回True。
3. 二分迭代：若检查通过，说明可以尝试更大的平均值，令left=mid；否则令right=mid。迭代约100次后，left即为所求的最大平均值。

代码实现 #

def is_possible(mid, arr, n, k):
    prefix = [0.0] * (n + 1)
    for i in range(1, n + 1):
        prefix[i] = prefix[i-1] + arr[i-1] - mid
    
    min_prefix = prefix[0]
    for i in range(k, n + 1):
        if prefix[i] >= min_prefix:
            return True
        # 更新前缀最小值，覆盖j <= i - k的范围
        min_prefix = min(min_prefix, prefix[i - k + 1])
    return False

n, k = map(int, input().split())
arr = list(map(int, input().split()))

left = 0.0
right = float(max(arr))
# 迭代100次保证精度足够
for _ in range(100):
    mid = (left + right) / 2
    if is_possible(mid, arr, n, k):
        left = mid
    else:
        right = mid

print("{0:.6f}".format(left))

时间复杂度 #

二分迭代次数为O(log(max_a))（max_a≤1e6，约20次），每次检查耗时O(n)，总时间复杂度为O(n log C)（C为数组元素最大值），完全适配n=3e5的数据规模。

内容的提问来源于stack exchange，提问作者Phạm Hồ Hữu Trí