LeetCode 813 - Largest Sum of Averages

The problem gives us an integer array nums and an integer k. We are allowed to split the array into at most k contiguous, non-empty subarrays. For each subarray, we compute its average, then sum all of those averages together. Our goal is to maximize that total score.

leetcodemediumarraydynamic-programmingprefix-sum

LeetCode Problem 813

Difficulty: 🟡 Medium
Topics: Array, Dynamic Programming, Prefix Sum

Solution

Problem Understanding

The problem gives us an integer array nums and an integer k. We are allowed to split the array into at most k contiguous, non-empty subarrays. For each subarray, we compute its average, then sum all of those averages together. Our goal is to maximize that total score.

The important detail is that the partitions must be adjacent and must cover the entire array. We cannot skip elements, reorder values, or create empty groups. Every element must belong to exactly one partition.

For example, if:

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

one valid partitioning is:

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

The score becomes:

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

The problem asks us to find the maximum possible score among all valid partitionings.

The constraints are relatively small:

  • 1 <= nums.length <= 100
  • 1 <= nums[i] <= 10^4
  • 1 <= k <= nums.length

Since the array size is only 100, this strongly suggests that a polynomial-time dynamic programming solution is appropriate. An exponential brute-force search would still be too slow because the number of possible partitions grows rapidly.

Several edge cases are important:

  • When k = 1, we cannot partition at all, so the answer is simply the average of the entire array.
  • When k = nums.length, every element can become its own subarray. Since the average of a single element is the element itself, the answer becomes the sum of all elements.
  • Arrays with repeated values can produce many equivalent partitionings.
  • Small arrays of length 1 must still work correctly.
  • Since averages are fractional, we must use floating-point arithmetic carefully.

Approaches

Brute Force Approach

A brute-force solution would try every possible way to partition the array into up to k groups.

At every index, we can either continue the current subarray or start a new partition. Once a partitioning is complete, we compute the average of every segment and keep track of the maximum score.

This approach is correct because it explores all possible valid partitionings, guaranteeing that the optimal answer will eventually be found.

However, the number of possible partitions grows exponentially with the size of the array. Even with only 100 elements, this becomes computationally infeasible.

Additionally, repeatedly computing subarray averages would introduce unnecessary repeated work unless prefix sums are used.

Dynamic Programming Insight

The key observation is that the problem has optimal substructure.

Suppose we decide that the last partition begins at index j and ends at index i. Then:

  • The score contributed by this last partition is the average of nums[j:i+1]
  • Everything before j must already be partitioned optimally into fewer groups

This naturally leads to dynamic programming.

We define:

dp[p][i]

as the maximum score obtainable by partitioning the first i elements into exactly p groups.

To transition:

dp[p][i] = max(
    dp[p-1][j] + average(nums[j:i])
)

for all valid j.

To compute averages efficiently, we use a prefix sum array.

This reduces the complexity from exponential to polynomial time.

Approach Comparison

Approach Time Complexity Space Complexity Notes
Brute Force O(2^n) O(n) Tries every partition configuration
Optimal Dynamic Programming O(k * n^2) O(k * n) Uses DP with prefix sums for efficient transitions

Algorithm Walkthrough

  1. Create a prefix sum array so that any subarray sum can be computed in constant time.

We define:

prefix[i]

as the sum of the first i elements.

Then the sum of nums[l:r] becomes:

prefix[r] - prefix[l]

This allows us to compute averages efficiently during DP transitions. 2. Define the DP state.

Let:

dp[p][i]

represent the maximum score obtainable by partitioning the first i elements into exactly p groups.

Here:

  • p ranges from 1 to k
  • i ranges from 1 to n
  1. Initialize the base case.

When there is only one partition:

dp[1][i] = average(nums[0:i])

Since the entire prefix must belong to one group. 4. Fill the DP table for additional partitions.

For every number of partitions p from 2 to k:

  • Consider every ending position i
  • Try every possible starting point j for the last partition

The recurrence becomes:

dp[p][i] = max(
    dp[p-1][j] + average(nums[j:i])
)
  1. Compute subarray averages using prefix sums.

The average of nums[j:i] is:

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

This avoids recomputing sums repeatedly. 6. Return the final answer.

Since we may use at most k partitions, the optimal solution will always benefit from additional partitions because all numbers are positive.

Therefore:

dp[k][n]

is the final answer.

Why it works

The algorithm works because every optimal partitioning can be decomposed into two parts:

  • An optimal partitioning of the prefix
  • One final subarray

The DP recurrence exhaustively considers every valid position where the final partition could begin. Since each smaller subproblem is solved optimally before being reused, the final result is guaranteed to be optimal by dynamic programming principles.

Python Solution

from typing import List

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

        # Prefix sums
        prefix = [0] * (n + 1)

        for i in range(n):
            prefix[i + 1] = prefix[i] + nums[i]

        # dp[p][i] = maximum score using first i elements with p groups
        dp = [[0.0] * (n + 1) for _ in range(k + 1)]

        # Base case: one partition
        for i in range(1, n + 1):
            dp[1][i] = prefix[i] / i

        # Fill DP table
        for p in range(2, k + 1):
            for i in range(p, n + 1):
                for j in range(p - 1, i):
                    current_average = (
                        prefix[i] - prefix[j]
                    ) / (i - j)

                    dp[p][i] = max(
                        dp[p][i],
                        dp[p - 1][j] + current_average
                    )

        return dp[k][n]

The implementation begins by constructing a prefix sum array. This allows constant-time computation of any subarray sum, which is essential because averages are repeatedly evaluated during DP transitions.

Next, the DP table is initialized. Each entry dp[p][i] stores the best score achievable using the first i elements partitioned into exactly p groups.

The base case handles the situation where only one partition is allowed. In that case, the score is simply the average of the entire prefix.

The nested loops then fill the DP table incrementally. For each possible number of partitions and each array prefix, the algorithm tries every possible split point for the last partition.

The transition computes:

  • The best score for the earlier prefix
  • The average of the final partition

and combines them.

Finally, the algorithm returns dp[k][n], which represents the optimal score using all elements and up to k partitions.

Go Solution

func largestSumOfAverages(nums []int, k int) float64 {
	n := len(nums)

	// Prefix sums
	prefix := make([]float64, n+1)

	for i := 0; i < n; i++ {
		prefix[i+1] = prefix[i] + float64(nums[i])
	}

	// dp[p][i] = max score using first i elements with p groups
	dp := make([][]float64, k+1)

	for i := range dp {
		dp[i] = make([]float64, n+1)
	}

	// Base case
	for i := 1; i <= n; i++ {
		dp[1][i] = prefix[i] / float64(i)
	}

	// Fill DP table
	for p := 2; p <= k; p++ {
		for i := p; i <= n; i++ {
			for j := p - 1; j < i; j++ {
				currentAverage := (prefix[i] - prefix[j]) / float64(i-j)

				candidate := dp[p-1][j] + currentAverage

				if candidate > dp[p][i] {
					dp[p][i] = candidate
				}
			}
		}
	}

	return dp[k][n]
}

The Go implementation closely mirrors the Python version. Since Go distinguishes between integers and floating-point values strictly, the prefix sum array is stored as float64 from the beginning to simplify average calculations.

Slices are used for the DP table, and all arithmetic involving averages uses explicit float64 conversion.

Unlike Python, Go does not have built-in max for floating-point numbers, so the implementation uses a manual comparison.

Worked Examples

Example 1

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

Step 1: Prefix Sums

Index Prefix Sum
0 0
1 9
2 10
3 12
4 15
5 24

Step 2: Base Case, One Partition

i Subarray Average dp[1][i]
1 [9] 9 9
2 [9,1] 5 5
3 [9,1,2] 4 4
4 [9,1,2,3] 3.75 3.75
5 [9,1,2,3,9] 4.8 4.8

Step 3: Two Partitions

For dp[2][5], try all split points:

j Left Score Right Partition Right Average Total
1 9 [1,2,3,9] 3.75 12.75
2 5 [2,3,9] 4.67 9.67
3 4 [3,9] 6 10
4 3.75 [9] 9 12.75

Best value:

dp[2][5] = 12.75

Step 4: Three Partitions

For dp[3][5]:

j dp[2][j] Last Partition Average Total
2 10 [2,3,9] 4.67 14.67
3 11 [3,9] 6 17
4 11 [9] 9 20

Final answer:

20

Example 2

nums = [1,2,3,4,5,6,7]
k = 4

The optimal partitioning becomes:

[1,2,3,4] [5] [6] [7]

The score is:

2.5 + 5 + 6 + 7 = 20.5

The DP systematically evaluates every valid split arrangement and eventually discovers this optimal partitioning.

Complexity Analysis

Measure Complexity Explanation
Time O(k * n^2) For each partition count and ending index, all split points are explored
Space O(k * n) DP table stores states for every partition count and prefix length

The outer loop iterates over partition counts from 1 to k. The middle loop iterates over array positions, and the inner loop tries every possible split point. This creates the O(k * n^2) runtime.

The DP table stores (k + 1) * (n + 1) floating-point values, resulting in O(k * n) space usage.

Test Cases

from typing import List

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

        prefix = [0] * (n + 1)

        for i in range(n):
            prefix[i + 1] = prefix[i] + nums[i]

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

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

        for p in range(2, k + 1):
            for i in range(p, n + 1):
                for j in range(p - 1, i):
                    avg = (prefix[i] - prefix[j]) / (i - j)

                    dp[p][i] = max(
                        dp[p][i],
                        dp[p - 1][j] + avg
                    )

        return dp[k][n]

solution = Solution()

# Provided example 1
assert abs(solution.largestSumOfAverages([9,1,2,3,9], 3) - 20.0) < 1e-6

# Provided example 2
assert abs(solution.largestSumOfAverages([1,2,3,4,5,6,7], 4) - 20.5) < 1e-6

# Single element array
assert abs(solution.largestSumOfAverages([5], 1) - 5.0) < 1e-6

# k = 1, entire array must stay together
assert abs(solution.largestSumOfAverages([1,2,3,4], 1) - 2.5) < 1e-6

# k = n, every element becomes its own partition
assert abs(solution.largestSumOfAverages([1,2,3,4], 4) - 10.0) < 1e-6

# Repeated values
assert abs(solution.largestSumOfAverages([5,5,5,5], 2) - 10.0) < 1e-6

# Increasing sequence
assert abs(solution.largestSumOfAverages([1,2,3,4,5], 2) - 7.5) < 1e-6

# Large value differences
assert abs(solution.largestSumOfAverages([100,1,1,1,100], 2) - 125.75) < 1e-6

# Many partitions allowed
assert abs(solution.largestSumOfAverages([10,20,30], 3) - 60.0) < 1e-6

# Minimal partitions with larger array
assert abs(solution.largestSumOfAverages([3,3,3,3,3], 1) - 3.0) < 1e-6
Test Why
[9,1,2,3,9], k=3 Validates the main example
[1,2,3,4,5,6,7], k=4 Tests multiple partition possibilities
[5], k=1 Smallest possible input
[1,2,3,4], k=1 Ensures no partitioning works correctly
[1,2,3,4], k=4 Ensures every element can become its own group
[5,5,5,5], k=2 Tests repeated identical values
[1,2,3,4,5], k=2 Tests increasing values
[100,1,1,1,100], k=2 Tests uneven distributions
[10,20,30], k=3 Tests maximal partition usage
[3,3,3,3,3], k=1 Tests stable averages

Edge Cases

One important edge case occurs when k = 1. In this scenario, no partitioning is allowed, so the algorithm must return the average of the entire array. A buggy implementation might still attempt transitions for additional groups or mishandle initialization. The base case initialization correctly handles this by directly assigning the average of every prefix to dp[1][i].

Another important case occurs when k = nums.length. Here, every element can become its own partition, which always maximizes the score because all numbers are positive. The implementation naturally handles this because the DP explores all valid split points, including partitions of size one.

Arrays of length one are also easy to mishandle. Since there is only one valid partition, the algorithm must return that element itself. The implementation works correctly because the prefix sum and DP initialization both support arrays with minimal size.

A more subtle edge case involves repeated values. When many partitions produce identical scores, floating-point comparisons can become tricky. Since the implementation uses standard floating-point arithmetic and only requires precision within 10^-6, it safely handles these cases.

Finally, arrays containing large values mixed with very small values can expose incorrect averaging logic or integer division bugs. Both implementations explicitly use floating-point division, ensuring accurate averages regardless of value distribution.