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.

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 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:

1. Base Case 1: When i == n, we've reached the end of the array, so we return 0 (no more elements to partition).

2. 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).

3. 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.

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:

1. First subarray = [4], remaining = [1, 7]
2. First subarray = [4, 1], remaining = [7]
3. 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.

Solution Implementation

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)