LeetCode 644: Maximum Average Subarray II

Problem Restatement

We are given an integer array nums and an integer k.

We need to find the maximum possible average value among all contiguous subarrays whose length is at least k.

The answer is accepted if the error is less than 10^-5. The constraints are 1 <= k <= n <= 10^4, and each nums[i] is between -10^4 and 10^4.

Input and Output

Example function shape:

def findMaxAverage(nums: list[int], k: int) -> float:
    ...

Examples

Example 1:

nums = [1, 12, -5, -6, 50, 3]
k = 4

Valid subarrays must have length 4, 5, or 6.

For length 4:

For length 5:

For length 6:

The maximum is:

12.75

Example 2:

nums = [5]
k = 1

The only subarray is [5], so the answer is:

5.0

First Thought: Try Every Subarray

A direct solution is to try every starting index and every ending index.

For each subarray with length at least k, compute its average and keep the maximum.

This is too slow because there are O(n^2) subarrays.

Even with prefix sums, checking all possible subarrays still takes O(n^2) time.

We need a different view.

Key Insight

Instead of directly finding the maximum average, ask a yes-or-no question:

Is there a subarray of length at least k with average at least x?

If the answer is yes, then the true maximum average is at least x.

If the answer is no, then the true maximum average is less than x.

That gives a monotonic condition, so we can binary search the answer.

The possible average lies between:

min(nums)

and:

max(nums)

Turning Average Into Sum

For a guessed average x, a subarray has average at least x if:

sum(subarray) / length >= x

This is equivalent to:

sum(subarray) >= x * length

Move the right side into the sum:

sum(num - x for num in subarray) >= 0

So for each guess x, subtract x from every number.

Now the question becomes:

Is there a subarray of length at least k whose transformed sum is non-negative?

This can be checked in linear time with prefix sums.

Feasibility Check

Let the transformed prefix sum be:

prefix[i] = sum(nums[0:i] - x)

For a subarray from j to i - 1, the transformed sum is:

prefix[i] - prefix[j]

The length is:

i - j

We need i - j >= k.

For each i, the best choice is the smallest prefix sum among all valid j <= i - k.

So while scanning, we maintain:

min_prefix = min(prefix[0], prefix[1], ..., prefix[i - k])

If:

prefix[i] - min_prefix >= 0

then there exists a valid subarray with average at least x.

Algorithm

1. Set left = min(nums) and right = max(nums).
2. Binary search while right - left > 1e-5.
3. Let mid = (left + right) / 2.
4. Check whether there is a subarray of length at least k with average at least mid.
5. If yes, set left = mid.
6. Otherwise, set right = mid.
7. Return left.

Implementation

class Solution:
    def findMaxAverage(self, nums: list[int], k: int) -> float:
        def can_find(average: float) -> bool:
            prefix = 0.0
            previous_prefix = 0.0
            min_previous_prefix = 0.0

            for i in range(k):
                prefix += nums[i] - average

            if prefix >= 0:
                return True

            for i in range(k, len(nums)):
                prefix += nums[i] - average
                previous_prefix += nums[i - k] - average
                min_previous_prefix = min(min_previous_prefix, previous_prefix)

                if prefix - min_previous_prefix >= 0:
                    return True

            return False

        left = min(nums)
        right = max(nums)

        while right - left > 1e-5:
            mid = (left + right) / 2

            if can_find(mid):
                left = mid
            else:
                right = mid

        return left

Code Explanation

The helper function answers:

can_find(average)

It returns True if some subarray of length at least k has average at least average.

First, we compute the transformed sum of the first k elements:

for i in range(k):
    prefix += nums[i] - average

If this sum is already non-negative, then the first window works:

if prefix >= 0:
    return True

Then we extend the right end one index at a time.

prefix += nums[i] - average

This is the transformed prefix sum up to the current position.

We also update the prefix sum that is exactly k positions behind:

previous_prefix += nums[i - k] - average

Then we keep the smallest such prefix:

min_previous_prefix = min(min_previous_prefix, previous_prefix)

If removing the smallest valid prefix leaves a non-negative sum:

if prefix - min_previous_prefix >= 0:
    return True

then a valid subarray exists.

The binary search moves upward when a candidate average is feasible:

if can_find(mid):
    left = mid

Otherwise, it moves downward:

else:
    right = mid

Correctness

For a candidate average x, a subarray has average at least x exactly when the sum of num - x over that subarray is non-negative.

The helper can_find(x) checks every possible right endpoint. For each right endpoint, it maintains the smallest transformed prefix sum among all left endpoints that make the subarray length at least k. Subtracting the smallest such prefix gives the largest transformed sum among valid subarrays ending there.

If this value is ever non-negative, then a subarray of length at least k has average at least x, so the helper returns True.

If the helper finishes without finding such a value, then no valid subarray has average at least x.

This feasibility condition is monotonic. If some average x is feasible, then every smaller average is also feasible. If x is not feasible, then every larger average is also not feasible.

Therefore, binary search over the average range converges to the maximum feasible average. Since the loop stops when the range is at most 1e-5, the returned value satisfies the required precision.

Complexity

Let n = len(nums) and let R = max(nums) - min(nums).

For this problem, eps = 1e-5.

Testing

def run_tests():
    s = Solution()

    assert abs(s.findMaxAverage([1, 12, -5, -6, 50, 3], 4) - 12.75) < 1e-5
    assert abs(s.findMaxAverage([5], 1) - 5.0) < 1e-5
    assert abs(s.findMaxAverage([-1], 1) - (-1.0)) < 1e-5
    assert abs(s.findMaxAverage([0, 4, 0, 3, 2], 1) - 4.0) < 1e-5
    assert abs(s.findMaxAverage([-5, -2, -3, -4], 2) - (-2.5)) < 1e-5
    assert abs(s.findMaxAverage([10, 10, 10], 2) - 10.0) < 1e-5

    print("all tests passed")

run_tests()

Test meaning: