Largest Sum of Averages

medium dp prefix-sum partition

Problem

Given an array nums and an integer k, partition nums into at most k non-empty contiguous groups. Return the maximum possible sum of group averages.

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

Output20.0

Cut into [9],[1,2,3],[9] gives 9+2+9 = 20.

nums, k

def largestSumOfAverages(nums, k):
    n = len(nums)
    pre = [0.0] * (n + 1)
    for i, v in enumerate(nums): pre[i+1] = pre[i] + v
    dp = [pre[n] - pre[i] for i in range(n + 1)]
    dp[n] = 0.0
    for groups in range(2, k + 1):
        new = dp[:]
        for i in range(n):
            best = 0.0
            for j in range(i + 1, n):
                avg = (pre[j] - pre[i]) / (j - i)
                if avg + dp[j] > best: best = avg + dp[j]
            new[i] = best
        dp = new
    return dp[0]

var largestSumOfAverages = function(nums, k) {
    const n = nums.length;
    const pre = new Array(n + 1).fill(0);
    for (let i = 0; i < n; i++) pre[i+1] = pre[i] + nums[i];
    let dp = new Array(n + 1).fill(0);
    for (let i = 0; i < n; i++) dp[i] = (pre[n] - pre[i]) / (n - i);
    for (let g = 2; g <= k; g++) {
        const nw = dp.slice();
        for (let i = 0; i < n; i++) {
            let best = 0;
            for (let j = i+1; j < n; j++) {
                const avg = (pre[j] - pre[i]) / (j - i);
                if (avg + dp[j] > best) best = avg + dp[j];
            }
            nw[i] = best;
        }
        dp = nw;
    }
    return dp[0];
};

class Solution {
    public double largestSumOfAverages(int[] nums, int k) {
        int n = nums.length;
        double[] pre = new double[n + 1];
        for (int i = 0; i < n; i++) pre[i+1] = pre[i] + nums[i];
        double[] dp = new double[n + 1];
        for (int i = 0; i < n; i++) dp[i] = (pre[n] - pre[i]) / (n - i);
        for (int g = 2; g <= k; g++) {
            double[] nw = dp.clone();
            for (int i = 0; i < n; i++) {
                double best = 0;
                for (int j = i+1; j < n; j++) {
                    double avg = (pre[j] - pre[i]) / (j - i);
                    if (avg + dp[j] > best) best = avg + dp[j];
                }
                nw[i] = best;
            }
            dp = nw;
        }
        return dp[0];
    }
}

class Solution {
public:
    double largestSumOfAverages(vector<int>& nums, int k) {
        int n = nums.size();
        vector<double> pre(n + 1, 0);
        for (int i = 0; i < n; i++) pre[i+1] = pre[i] + nums[i];
        vector<double> dp(n + 1, 0);
        for (int i = 0; i < n; i++) dp[i] = (pre[n] - pre[i]) / (n - i);
        for (int g = 2; g <= k; g++) {
            vector<double> nw = dp;
            for (int i = 0; i < n; i++) {
                double best = 0;
                for (int j = i+1; j < n; j++) {
                    double avg = (pre[j] - pre[i]) / (j - i);
                    if (avg + dp[j] > best) best = avg + dp[j];
                }
                nw[i] = best;
            }
            dp = nw;
        }
        return dp[0];
    }
};

Explanation

We cut the array into at most k contiguous groups and add up each group's average; we want the largest total. The trick is to decide cuts with dynamic programming and use prefix sums so any group's average is instant.

First we build pre, the running prefix sums, so the average of nums[i..j) is just (pre[j] - pre[i]) / (j - i) with no extra looping.

Let dp[i] be the best total for the suffix starting at index i. With only one group allowed, dp[i] is simply the average of everything from i to the end.

Then we increase the group budget from 2 up to k. For each start i we try every spot j to end the first group, taking avg(i..j) + dp[j] and keeping the best. After processing all budgets, dp[0] is the answer.

Example: nums = [9,1,2,3,9], k = 3. The best cut is [9],[1,2,3],[9] giving 9 + 2 + 9 = 20.

Time: O(k * n^2) Space: O(n)