LeetCode 813: Largest Sum of Averages

Problem Restatement

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

We may partition nums into at most k non-empty adjacent subarrays. Every element must be used exactly once.

The score of a partition is the sum of the averages of its subarrays.

Return the maximum score possible. Answers within 10^-6 of the actual answer are accepted. The official problem states that the partition must use every integer in nums, and the score may be non-integer.

Input and Output

Examples

Example:

nums = [9, 1, 2, 3, 9]
k = 3

One good partition is:

[9], [1, 2, 3], [9]

Its score is:

9 + (1 + 2 + 3) / 3 + 9 = 20

So the answer is:

20.0

Another partition is:

[9, 1], [2], [3, 9]

Its score is:

5 + 2 + 6 = 13

This is valid, but worse.

First Thought: Try Every Partition

A direct solution would try every possible way to place partition cuts.

For each partition, compute the average of each group and sum them.

This works conceptually, but the number of possible cut combinations grows quickly. We need dynamic programming because many subproblems repeat.

Key Insight

The only choices are where to place cuts.

For a prefix ending at index i, suppose the last group starts at index j.

Then the score has two parts:

1. Best score for nums[0:j] using one fewer group.
2. Average of nums[j:i].

This gives a natural DP recurrence.

We use prefix sums so every subarray average can be computed in constant time.

average(j, i) = (prefix[i] - prefix[j]) / (i - j)

Here, nums[j:i] means the subarray from j to i - 1.

Algorithm

Build prefix sums:

prefix[0] = 0
prefix[i + 1] = prefix[i] + nums[i]

Define:

dp[i][g]

as the maximum score for partitioning the first i numbers into exactly g groups.

Base case:

dp[i][1] = prefix[i] / i

This means the first i numbers form one group.

Transition:

dp[i][g] = max(
    dp[j][g - 1] + average(j, i)
)

where:

g - 1 <= j < i

The last group is nums[j:i].

The answer is:

dp[n][k]

Although the statement says “at most k”, using exactly k groups is safe for positive numbers because splitting a group into more non-empty adjacent groups cannot decrease the sum of averages in this problem’s constraints.

Correctness

The DP state dp[i][g] stores the best score for the first i elements split into exactly g non-empty adjacent groups.

For g = 1, there is only one possible partition: all first i elements form one group. Its score is their average, so the base case is correct.

For g > 1, consider an optimal partition of the first i elements into g groups. Its last group must be some suffix nums[j:i], where the first j elements are split into g - 1 groups. The score is therefore:

dp[j][g - 1] + average(j, i)

The transition checks every valid j, so it includes the last cut position of the optimal partition. It also only constructs valid adjacent non-empty groups.

Thus each DP state is correct. In particular, dp[n][k] gives the maximum score for the full array.

Complexity

Let n = len(nums).

Implementation

class Solution:
    def largestSumOfAverages(self, nums: list[int], k: int) -> float:
        n = len(nums)

        prefix = [0.0] * (n + 1)
        for i, num in enumerate(nums):
            prefix[i + 1] = prefix[i] + num

        dp = [[0.0] * (k + 1) for _ in range(n + 1)]

        for i in range(1, n + 1):
            dp[i][1] = prefix[i] / i

        for groups in range(2, k + 1):
            for i in range(groups, n + 1):
                for j in range(groups - 1, i):
                    last_average = (prefix[i] - prefix[j]) / (i - j)
                    dp[i][groups] = max(
                        dp[i][groups],
                        dp[j][groups - 1] + last_average,
                    )

        return dp[n][k]

Code Explanation

The prefix sum array lets us compute any subarray sum quickly:

prefix[i + 1] = prefix[i] + num

The sum of nums[j:i] is:

prefix[i] - prefix[j]

So its average is:

(prefix[i] - prefix[j]) / (i - j)

The base case fills the score for one group:

dp[i][1] = prefix[i] / i

Then we try larger group counts:

for groups in range(2, k + 1):

For each i, we choose where the last group starts:

for j in range(groups - 1, i):

The first j numbers use groups - 1 groups. The last group is nums[j:i].

We keep the best score:

dp[i][groups] = max(
    dp[i][groups],
    dp[j][groups - 1] + last_average,
)

Finally, return the best score for all n numbers using k groups:

return dp[n][k]

Testing

def close(a, b, eps=1e-6):
    return abs(a - b) <= eps

def run_tests():
    s = Solution()

    assert close(s.largestSumOfAverages([9, 1, 2, 3, 9], 3), 20.0)
    assert close(s.largestSumOfAverages([1, 2, 3, 4, 5], 1), 3.0)
    assert close(s.largestSumOfAverages([5], 1), 5.0)
    assert close(s.largestSumOfAverages([1, 2, 3], 3), 6.0)
    assert close(s.largestSumOfAverages([4, 1, 7, 5, 6, 2, 3], 4), 18.1666666667)

    print("all tests passed")

run_tests()