「刷题」Maximum Average Subarray II

Link: leetcode644 (locked), lintcode617

Description

给出一个整数数组，有正有负。找到这样一个连续的子数组，他的长度大于等于 kk，且平均值最大。

Solution

连续部分求和，想到用前缀和处理，S[i]=∑t=0inums[t]S[i] = \sum_{t=0}^{i}{\text{nums}[t]}.

原来的问题就被转化成：有 nn 个离散点，坐标分别为：(1,S[1]),⋯ ,(n,S[n])(1,S[1]), \cdots, (n, S[n])，最大化 S[j]−S[i]j−i\displaystyle\frac{S[j]-S[i]}{j-i}, 并满足 j−i≥kj-i \ge k.

回到自然语言描述，这是一个计算几何问题：平面上的点集，找一条线段两端点横坐标差至少为 kk, 使得斜率最大。

维护一个 deque；从k开始枚举右边端点的横坐标 j， 对于一个确定的 j：

  • 判断最近的临界点 i=j-k 是否能加入 deque 的队列尾，注意维护deque中的点集构成一个相邻点之间斜率递增的下凸包 (lower hull)
  • 从deque的队列头开始，贪心地选择斜率最大的点作为切点即可（只要队头的点不与j形成最大斜率,就不断popleft），有点像凸优化的梯度下降

Code

import collections
from typing import (
    List,
)


class Solution:
    """
    @param nums: an array with positive and negative numbers
    @param k: an integer
    @return: the maximum average
    """
    def max_average(self, nums: List[int], k: int) -> float:
        n = len(nums)
        s = [0] * (n + 1)
        for i in range(n):
            s[i+1] = s[i] + nums[i]


        max_avg = -float('inf')
        deque = collections.deque()


        def grad(x1: int, x2: int) -> float:
            return (s[x2]-s[x1]) / (x2-x1)


        # Enumerate j, where j - i >= k.
        for j in range(k, n + 1):
            i = j - k
            # Maintain the lower hull
            while len(deque) >= 2 and grad(deque[-1], i) <= grad(deque[-2], deque[-1]):
                deque.pop()
            deque.append(i)
            # Greedy
            while len(deque) >= 2 and grad(deque[0], j) <= grad(deque[1], j):
                deque.popleft()
            max_avg = max(grad(deque[0], j), max_avg)
        return max_avg