题目描述

有一个整数数组，大小为$n$，数组中的元素均为正整数。现在定义“子数组”为原数组中连续的一段元素，并且需要至少有一个元素。请你编写一个时间复杂度为$O(nlogn)$的算法，找出数组中所有子数组中元素的平均值最大的子数组，并返回该子数组的起始位置和长度。如果有多个满足条件的子数组，返回其中任意一个即可。

输入格式

第一行一个整数$n$，表示数组的大小。

第二行为$n$个整数$a[i]$。

输出格式

输出两个整数，分别为最大平均值对应的子数组的起始位置和长度。如果有多个满足条件的子数组，输出其中任意一个即可。

输入样例

6
1 12 -5 -6 50 3

输出样例

4 2

解题思路

题目求出数组中所有子数组的平均值最大的子数组的起始位置和长度。我们可以用二分+前缀和来思考。

我们可以从0到数组长度之间二分答案mid，假设当前选取的值为mid，并且我们已经得到了所有平均值大于等于mid的子数组。

首先，我们将数组中的每个数减去mid，然后计算其前缀和。这样我们可以得到一个新数组$sum[i]$，表示原数组中前i个数减去mid后的前缀和。然后对$sum[i - 1]$做单调栈，我们可以求出其中的最小值和其对应的下标$j$。接着，我们可以推出以i为结尾且平均值大于等于mid的所有子数组中，以j + 1为起点的子数组的最大平均值即$sum[i] - sum[j] \geq (i - j) \times mid$，即$sum[i] - sum[j] \geq i \times mid - j \times mid$，即$\frac{sum[i] - i \times mid}{sum[j] - j \times mid} \geq 1$。

同时，我们可以保证$sum[j]$是当前数组中所有元素减去mid形成的数组中前j个元素的前缀和，即在区间$[0, j - 1]$中，没有以平均值大于等于mid的子数组。因此，我们只需要判断$sum[j]$是否小于等于0，如果成立，那么我们可以得到一个以i为结尾且平均值大于等于mid的子数组，它的起点必须在j+1之后。

由于数组中可能有负数，我们需要先对$sum$数组做一次前缀和，然后对$sum[i - 1]$做一次单调栈。由于单调栈中的每个元素只会被push一次和pop一次，所以总时间复杂度为$O(n)$。二分的过程中，每次执行一次减法和一次单调栈操作，时间复杂度为$O(n)$。因此，总时间复杂度为$O(nlogn)$。

代码实现

def find_max_avg_subarray(n, a):
    s = [0] * (n + 1)
    for i in range(n):
        s[i + 1] = s[i] + a[i]
    l, r = 0, max(a)
    while l < r:
        mid = (l + r + 1) // 2
        max_sum, max_j = float('-inf'), 0
        for i in range(1, n + 1):
            j = max_j
            while j < i and s[i] - s[j] - mid * (i - j) >= 0:
                j += 1
            if s[i] - s[j] >= mid * (i - j):
                max_sum, max_j = s[i] - s[j], j
        if max_sum >= 0:
            l = mid
        else:
            r = mid - 1
    for i in range(1, n + 1):
        j = max_j
        while j < i and s[i] - s[j] - l * (i - j) >= 0:
            j += 1
        if s[i] - s[j] >= l * (i - j):
            return j, i - j