813. Largest Sum of Averages

MediumArrayDynamic ProgrammingPrefix Sum

Problem Description

You have an integer array nums and an integer k. Your task is to partition this array into at most k non-empty adjacent subarrays to maximize a score.

The score of a partition is calculated as the sum of the averages of each subarray in the partition.

Key requirements:

You must use every element in nums (the partition must be complete)
The subarrays must be adjacent (consecutive elements)
You can use at most k subarrays (you can use fewer than k)
The score doesn't need to be an integer (it can be a decimal)

Your goal is to find the maximum possible score among all valid partitions.

For example, if you have nums = [9, 1, 2, 3, 9] and k = 3, one possible partition could be [9], [1, 2, 3], [9] with score: 9/1 + (1+2+3)/3 + 9/1 = 9 + 2 + 9 = 20. Another partition could be [9, 1], [2, 3, 9] with score: (9+1)/2 + (2+3+9)/3 = 5 + 4.667 = 9.667. You need to find the partition that gives the maximum score.

The answer will be accepted if it's within 10^-6 of the actual answer.

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

To maximize the sum of averages, we need to understand a key mathematical property: dividing an array into more groups generally increases the total score. This is because the average of averages tends to be higher than the average of the whole. For instance, [9]/1 + [1]/1 = 10 is greater than [9,1]/2 = 5.

This suggests we should use as many partitions as possible (up to k) to maximize our score. However, we can't just split arbitrarily - we need to be strategic about where we make the cuts.

The problem has an optimal substructure property: if we decide to make the first subarray from index 0 to some index j-1, then the remaining problem becomes finding the best way to partition nums[j:] into at most k-1 groups. This naturally leads to a dynamic programming approach.

We can think of this recursively: at each position i, we try all possible endpoints for the current subarray (from i+1 to n), calculate the average of that subarray, and add it to the best possible score for partitioning the remaining array with one fewer group allowed.

To efficiently calculate subarray sums, we can use a prefix sum array s where s[i] represents the sum of elements from index 0 to i-1. This allows us to compute the sum of any subarray from index i to j-1 as s[j] - s[i] in constant time.

The base cases are straightforward:

When we've processed all elements (i = n), the score is 0
When we have only one group left (k = 1), we must take all remaining elements as a single subarray

By using memoization with @cache, we avoid recalculating the same subproblems multiple times, making the solution efficient.

Learn more about Dynamic Programming and Prefix Sum patterns.

Solution Implementation

1class Solution: 2 def largestSumOfAverages(self, nums: List[int], k: int) -> float: 3 """ 4 Find the largest sum of averages after partitioning array into k groups. 5 6 Args: 7 nums: Input array of integers 8 k: Number of groups to partition the array into 9 10 Returns: 11 Maximum sum of averages possible 12 """ 13 from functools import cache 14 from itertools import accumulate 15 from typing import List 16 17 @cache 18 def dp(start_idx: int, groups_left: int) -> float: 19 """ 20 Dynamic programming function to find max sum of averages. 21 22 Args: 23 start_idx: Starting index of current subarray 24 groups_left: Number of groups remaining to form 25 26 Returns: 27 Maximum sum of averages from start_idx to end with groups_left partitions 28 """ 29 # Base case: reached end of array 30 if start_idx == array_length: 31 return 0 32 33 # Base case: only one group left, take average of remaining elements 34 if groups_left == 1: 35 remaining_sum = prefix_sum[array_length] - prefix_sum[start_idx] 36 remaining_count = array_length - start_idx 37 return remaining_sum / remaining_count 38 39 max_sum = 0 40 # Try all possible end positions for current group 41 for end_idx in range(start_idx + 1, array_length): 42 # Calculate average of current group [start_idx, end_idx) 43 current_group_sum = prefix_sum[end_idx] - prefix_sum[start_idx] 44 current_group_size = end_idx - start_idx 45 current_average = current_group_sum / current_group_size 46 47 # Recursively calculate maximum for remaining elements 48 remaining_max = dp(end_idx, groups_left - 1) 49 50 # Update maximum sum of averages 51 max_sum = max(max_sum, current_average + remaining_max) 52 53 return max_sum 54 55 # Initialize variables 56 array_length = len(nums) 57 # Create prefix sum array for efficient range sum calculation 58 prefix_sum = list(accumulate(nums, initial=0)) 59 60 # Start dynamic programming from index 0 with k groups 61 return dp(0, k) 62

1class Solution { 2 // Memoization table: dp[i][j] represents the maximum sum of averages 3 // starting from index i with j partitions remaining 4 private Double[][] memo; 5 6 // Prefix sum array for quick range sum calculation 7 private int[] prefixSum; 8 9 // Length of the input array 10 private int arrayLength; 11 12 public double largestSumOfAverages(int[] nums, int k) { 13 arrayLength = nums.length; 14 15 // Initialize prefix sum array with size n+1 16 prefixSum = new int[arrayLength + 1]; 17 18 // Initialize memoization table 19 memo = new Double[arrayLength][k + 1]; 20 21 // Build prefix sum array: prefixSum[i] = sum of nums[0] to nums[i-1] 22 for (int i = 0; i < arrayLength; i++) { 23 prefixSum[i + 1] = prefixSum[i] + nums[i]; 24 } 25 26 // Start DFS from index 0 with k partitions 27 return dfs(0, k); 28 } 29 30 /** 31 * Recursive function to find the maximum sum of averages 32 * @param startIndex - current starting index in the array 33 * @param remainingPartitions - number of partitions left to make 34 * @return maximum sum of averages from startIndex with remainingPartitions 35 */ 36 private double dfs(int startIndex, int remainingPartitions) { 37 // Base case: reached end of array 38 if (startIndex == arrayLength) { 39 return 0; 40 } 41 42 // Base case: only one partition left, take average of remaining elements 43 if (remainingPartitions == 1) { 44 int sumOfRemaining = prefixSum[arrayLength] - prefixSum[startIndex]; 45 int countOfRemaining = arrayLength - startIndex; 46 return (double) sumOfRemaining / countOfRemaining; 47 } 48 49 // Check if already computed 50 if (memo[startIndex][remainingPartitions] != null) { 51 return memo[startIndex][remainingPartitions]; 52 } 53 54 double maxSum = 0; 55 56 // Try all possible end points for the current partition 57 // endIndex represents the exclusive end of current partition 58 for (int endIndex = startIndex + 1; endIndex < arrayLength; endIndex++) { 59 // Calculate average of current partition [startIndex, endIndex) 60 int currentPartitionSum = prefixSum[endIndex] - prefixSum[startIndex]; 61 int currentPartitionSize = endIndex - startIndex; 62 double currentAverage = (double) currentPartitionSum / currentPartitionSize; 63 64 // Recursively calculate the best sum for remaining elements 65 double remainingSum = dfs(endIndex, remainingPartitions - 1); 66 67 // Update maximum sum 68 maxSum = Math.max(maxSum, currentAverage + remainingSum); 69 } 70 71 // Store result in memoization table and return 72 memo[startIndex][remainingPartitions] = maxSum; 73 return maxSum; 74 } 75} 76

1class Solution { 2public: 3 double largestSumOfAverages(vector<int>& nums, int k) { 4 int n = nums.size(); 5 6 // Prefix sum array to calculate range sums efficiently 7 // prefixSum[i] = sum of nums[0] to nums[i-1] 8 vector<int> prefixSum(n + 1, 0); 9 for (int i = 0; i < n; ++i) { 10 prefixSum[i + 1] = prefixSum[i] + nums[i]; 11 } 12 13 // Memoization table for dynamic programming 14 // memo[i][j] = maximum sum of averages starting from index i with j groups remaining 15 vector<vector<double>> memo(n, vector<double>(k + 1, -1.0)); 16 17 // Recursive function with memoization to find maximum sum of averages 18 // Parameters: 19 // startIdx: starting index of current subarray 20 // groupsLeft: number of groups we can still create 21 // Returns: maximum sum of averages from startIdx to end with groupsLeft groups 22 function<double(int, int)> dfs = [&](int startIdx, int groupsLeft) -> double { 23 // Base case: reached end of array 24 if (startIdx == n) { 25 return 0.0; 26 } 27 28 // Base case: only one group left, must take all remaining elements 29 if (groupsLeft == 1) { 30 int remainingSum = prefixSum[n] - prefixSum[startIdx]; 31 int remainingCount = n - startIdx; 32 return static_cast<double>(remainingSum) / remainingCount; 33 } 34 35 // Check if already computed 36 if (memo[startIdx][groupsLeft] >= 0) { 37 return memo[startIdx][groupsLeft]; 38 } 39 40 double maxSum = 0.0; 41 42 // Try all possible end positions for current group 43 // We need at least groupsLeft-1 elements after endIdx for remaining groups 44 for (int endIdx = startIdx + 1; endIdx <= n - groupsLeft + 1; ++endIdx) { 45 // Calculate average of current group [startIdx, endIdx) 46 int currentSum = prefixSum[endIdx] - prefixSum[startIdx]; 47 int currentCount = endIdx - startIdx; 48 double currentAverage = static_cast<double>(currentSum) / currentCount; 49 50 // Recursively calculate maximum for remaining elements 51 double remainingMax = dfs(endIdx, groupsLeft - 1); 52 53 // Update maximum sum 54 maxSum = max(maxSum, currentAverage + remainingMax); 55 } 56 57 // Store result in memoization table and return 58 return memo[startIdx][groupsLeft] = maxSum; 59 }; 60 61 // Start recursion from index 0 with k groups 62 return dfs(0, k); 63 } 64}; 65

1/** 2 * Finds the largest sum of averages after partitioning an array into k groups 3 * @param nums - The input array of numbers 4 * @param k - The number of groups to partition the array into 5 * @returns The maximum sum of averages possible 6 */ 7function largestSumOfAverages(nums: number[], k: number): number { 8 const arrayLength: number = nums.length; 9 10 // Build prefix sum array for quick range sum calculation 11 // prefixSum[i] represents sum of elements from index 0 to i-1 12 const prefixSum: number[] = Array(arrayLength + 1).fill(0); 13 for (let i = 0; i < arrayLength; i++) { 14 prefixSum[i + 1] = prefixSum[i] + nums[i]; 15 } 16 17 // Memoization table for dynamic programming 18 // memo[i][j] stores the maximum sum of averages starting from index i with j groups remaining 19 const memo: number[][] = Array.from( 20 { length: arrayLength }, 21 () => Array(k + 1).fill(0) 22 ); 23 24 /** 25 * Recursive function with memoization to find maximum sum of averages 26 * @param startIndex - Current starting index in the array 27 * @param groupsRemaining - Number of groups remaining to form 28 * @returns Maximum sum of averages from startIndex with groupsRemaining groups 29 */ 30 const calculateMaxAverage = (startIndex: number, groupsRemaining: number): number => { 31 // Base case: reached end of array 32 if (startIndex === arrayLength) { 33 return 0; 34 } 35 36 // Return memoized result if already calculated 37 if (memo[startIndex][groupsRemaining] > 0) { 38 return memo[startIndex][groupsRemaining]; 39 } 40 41 // If only one group remains, calculate average of all remaining elements 42 if (groupsRemaining === 1) { 43 const remainingSum: number = prefixSum[arrayLength] - prefixSum[startIndex]; 44 const remainingCount: number = arrayLength - startIndex; 45 return remainingSum / remainingCount; 46 } 47 48 // Try all possible end positions for the current group 49 for (let endIndex = startIndex + 1; endIndex < arrayLength; endIndex++) { 50 // Calculate average of current group from startIndex to endIndex-1 51 const currentGroupSum: number = prefixSum[endIndex] - prefixSum[startIndex]; 52 const currentGroupSize: number = endIndex - startIndex; 53 const currentGroupAverage: number = currentGroupSum / currentGroupSize; 54 55 // Recursively calculate maximum for remaining elements 56 const remainingMaxAverage: number = calculateMaxAverage(endIndex, groupsRemaining - 1); 57 58 // Update maximum sum of averages 59 memo[startIndex][groupsRemaining] = Math.max( 60 memo[startIndex][groupsRemaining], 61 currentGroupAverage + remainingMaxAverage 62 ); 63 } 64 65 return memo[startIndex][groupsRemaining]; 66 }; 67 68 // Start calculation from index 0 with k groups 69 return calculateMaxAverage(0, k); 70} 71

Solution Approach

The solution uses prefix sum and memoized search (dynamic programming) to efficiently find the maximum score.

Step 1: Build Prefix Sum Array

First, we create a prefix sum array s using accumulate(nums, initial=0). This gives us an array where s[i] represents the sum of all elements from index 0 to i-1. With this preprocessing, we can calculate the sum of any subarray from index i to j-1 as s[j] - s[i] in O(1) time.

Step 2: Define the Recursive Function

We define dfs(i, k) which returns the maximum sum of averages when partitioning the array starting from index i into at most k groups.

The function works as follows:

Base Case 1: When i == n, we've reached the end of the array, so we return 0 (no more elements to partition).
Base Case 2: When k == 1, we have only one group left, so we must take all remaining elements as a single subarray. The average is (s[n] - s[i]) / (n - i).
Recursive Case: For each possible endpoint j of the current subarray (from i+1 to n):
- Calculate the average of the current subarray: (s[j] - s[i]) / (j - i)
- Add the maximum score for partitioning the remaining array: dfs(j, k - 1)
- Track the maximum among all possible choices

Step 3: Memoization

The @cache decorator automatically memoizes the results of dfs(i, k), preventing redundant calculations. This is crucial because the same subproblem (i, k) may be encountered multiple times through different recursion paths.

Step 4: Return the Result

We call dfs(0, k) to start the partitioning from index 0 with at most k groups allowed.

Time Complexity: O(n² × k) where n is the length of the array, as we have n possible values for i, k possible values for k, and each state requires O(n) time to compute.

Space Complexity: O(n × k) for the memoization cache plus O(n) for the recursion stack.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through the solution with nums = [4, 1, 7] and k = 2.

Step 1: Build Prefix Sum Array

s = [0, 4, 5, 12] (cumulative sums: 0, 0+4, 0+4+1, 0+4+1+7)

Step 2: Calculate dfs(0, 2) (partition entire array with at most 2 groups)

We need to decide where to end the first subarray. We have three options:

First subarray = [4], remaining = [1, 7]
First subarray = [4, 1], remaining = [7]
First subarray = [4, 1, 7], remaining = []

Option 1: First subarray ends at index 1

Average of [4] = (s[1] - s[0]) / (1 - 0) = 4/1 = 4
Recursively solve dfs(1, 1) for [1, 7] with 1 group left
- With k=1, must take all remaining: average = (s[3] - s[1]) / (3 - 1) = 8/2 = 4
Total score: 4 + 4 = 8

Option 2: First subarray ends at index 2

Average of [4, 1] = (s[2] - s[0]) / (2 - 0) = 5/2 = 2.5
Recursively solve dfs(2, 1) for [7] with 1 group left
- With k=1, must take all remaining: average = (s[3] - s[2]) / (3 - 2) = 7/1 = 7
Total score: 2.5 + 7 = 9.5

Option 3: First subarray ends at index 3

Average of [4, 1, 7] = (s[3] - s[0]) / (3 - 0) = 12/3 = 4
Recursively solve dfs(3, 1) for [] (empty array)
- Base case: i=n, return 0
Total score: 4 + 0 = 4

Step 3: Choose Maximum

Maximum of (8, 9.5, 4) = 9.5

Therefore, the optimal partition is [4, 1], [7] with a maximum score of 9.5.

The memoization ensures that if we encounter the same (i, k) pair again (which doesn't happen in this small example), we reuse the computed result instead of recalculating.

Time and Space Complexity

The time complexity is O(n² × k) where n is the length of the array nums.

This complexity arises from the memoized recursive function dfs(i, k):

There are n × k possible states (combinations of starting index i from 0 to n-1, and partition count k from 1 to k)
For each state dfs(i, k), we iterate through positions j from i+1 to n-1, which takes O(n) time in the worst case
Therefore, the total time complexity is O(n × k) × O(n) = O(n² × k)

The space complexity is O(n × k).

This is due to:

The memoization cache stores up to n × k states from the recursive calls
The recursion stack depth is at most O(k) (we make at most k partitions)
The prefix sum array s takes O(n) space
Since k ≤ n, the dominant factor is the memoization cache, giving us O(n × k) space complexity

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Off-by-One Error in Loop Boundary

A critical pitfall occurs when iterating through possible endpoints for the current subarray. Many developers mistakenly write:

# INCORRECT - doesn't consider taking all remaining elements for end_idx in range(start_idx + 1, array_length):  # ... calculate average and recurse

This loop excludes the possibility of taking all remaining elements as a single group when groups_left > 1. The loop should actually go up to and include array_length:

# CORRECT - considers all valid partitions for end_idx in range(start_idx + 1, array_length + 1):  # ... calculate average and recurse

Why this matters: Without this fix, you miss valid partitions where the current group extends to the end of the array, potentially missing the optimal solution.

2. Integer Division Instead of Float Division

In Python 2 or when using floor division, you might accidentally use integer division:

# INCORRECT - loses precision current_average = (prefix_sum[end_idx] - prefix_sum[start_idx]) // (end_idx - start_idx)

Always ensure you're using float division:

# CORRECT - maintains precision current_average = (prefix_sum[end_idx] - prefix_sum[start_idx]) / (end_idx - start_idx)

3. Incorrect Prefix Sum Indexing

A common mistake is misunderstanding how prefix sums work:

# INCORRECT - wrong sum calculation current_sum = prefix_sum[end_idx - 1] - prefix_sum[start_idx]

The correct indexing for getting sum from index i to j-1 is:

# CORRECT - proper prefix sum usage current_sum = prefix_sum[j] - prefix_sum[i]

4. Missing Edge Case: k ≥ n

When k is greater than or equal to the array length, the optimal solution is to put each element in its own group (maximizing the sum since each element becomes its own average). Not handling this can lead to unnecessary computation:

# OPTIMIZATION - handle special case if k >= len(nums):  return sum(nums) # Each element in its own group

Complete Corrected Solution:

class Solution:  def largestSumOfAverages(self, nums: List[int], k: int) -> float:  from functools import cache  from itertools import accumulate   n = len(nums)   # Edge case optimization  if k >= n:  return float(sum(nums))   prefix_sum = list(accumulate(nums, initial=0))   @cache  def dp(i: int, k: int) -> float:  if i == n:  return 0   if k == 1:  return (prefix_sum[n] - prefix_sum[i]) / (n - i)   max_score = 0  # CRITICAL FIX: range should go to n+1, not n  for j in range(i + 1, n + 1):  avg = (prefix_sum[j] - prefix_sum[i]) / (j - i)  remaining = dp(j, k - 1)  max_score = max(max_score, avg + remaining)   return max_score   return dp(0, k)

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

Which of the following problems can be solved with backtracking (select multiple)

Generating subsets

Generating random numbers

Generating permutations

Sorting integers