Topics › Dynamic Programming › MCM DP - Partition DP

Partition Array for Maximum Sum

Medium TUF ID 810 A2Z

LeetCode ↗ Java source ↗ Edit this article ↗

Jump to code: Brute Better Optimal

Problem

You are given an integer array arr and an integer k. Split the array into contiguous pieces, each of length at most k. Then replace every element of each piece by the maximum of that piece. Return the largest possible sum of the array after this change.

Constraints (LeetCode 1043): 1 ≤ n ≤ 500, 0 ≤ arr[i] ≤ 10^9, 1 ≤ k ≤ n, and the answer fits in a 32-bit integer. The Java methods also accept an empty array (answer 0).

Example 1

arr = [1, 15, 7, 9, 2, 5, 10], k = 3
answer = 84              // [1, 15, 7] [9] [2, 5, 10] becomes [15, 15, 15, 9, 10, 10, 10]

Example 2

arr = [1, 4, 1, 5, 7, 3, 6, 1, 9, 9, 3], k = 4
answer = 83              // [1] [4, 1, 5, 7] [3, 6, 1, 9] [9, 3]: 1 + 28 + 36 + 18

Example 3

arr = [3, 8, 2], k = 1
answer = 13              // every piece has length 1, so nothing changes

Seeing the recurrence

A piece of length len whose largest value is mx contributes mx · len to the sum. Large values want to “spread” over as many neighbours as possible, but each piece is capped at k elements, and the pieces must tile the array, so choices interact.

Read the partition from left to right. The first piece starts at index i and has some length len between 1 and k (and it cannot run past the end). Its contribution depends only on its own elements. Everything after it, arr[i + len ..], is the same problem on a shorter suffix, and the best way to handle that suffix does not depend on how the first piece was chosen. That gives a one-index state:

f(n) = 0
f(i) = max over len in [1, min(k, n - i)] of  max(arr[i .. i+len-1]) · len + f(i + len)
answer = f(0)

This is the “front partition” pattern, the same shape as Palindrome Partitioning II. There, every piece had to pass a validity test (being a palindrome). Here every piece of length at most k is allowed, and the work is in the piece’s value.

One implementation detail matters for speed. As len grows by one, the piece gains exactly one element, arr[i + len - 1]. So the maximum can be updated with a running maxVal, in O(1) per step, instead of rescanning the piece.

Approach 1: Recursion

Intuition. Translate the recurrence directly: from position i, try every first-piece length, add its value to the best result for the rest, and keep the maximum.

Algorithm.

1. If i == n, return 0.
2. For len = 1 .. k while i + len ≤ n: update maxVal with arr[i + len - 1], candidate = maxVal · len + f(i + len).
3. Return the best candidate.

static int recursive(int[] arr, int k) {
    return solve(0, arr, k);
}

static int solve(int i, int[] arr, int k) {
    if (i == arr.length) return 0;
    int maxVal = 0, best = 0;
    for (int len = 1; len <= k && i + len <= arr.length; len++) {
        maxVal = Math.max(maxVal, arr[i + len - 1]);        // running maximum of arr[i .. i + len - 1]
        best = Math.max(best, maxVal * len + solve(i + len, arr, k));
    }
    return best;
}

Complexity. Time exponential: the recursion enumerates every way to cut the array into pieces of length at most k. When k ≥ n that is 2^(n-1) partitions; in general it is bounded by k^n. Space O(n) for the recursion depth.

Approach 2: Memoization

Intuition. Many different partitions of the prefix arrive at the same suffix start i. For example, pieces [a][b] and [a, b] both lead to f(2). There are only n + 1 distinct suffixes, so cache f(i) in dp[i]. Sums are never negative, so -1 marks an empty slot.

Algorithm.

1. Fill dp with -1.
2. Same recursion, returning dp[i] if known and storing each new result.

static int memoization(int[] arr, int k) {
    int[] dp = new int[arr.length];
    Arrays.fill(dp, -1);
    return memo(0, arr, k, dp);
}

static int memo(int i, int[] arr, int k, int[] dp) {
    if (i == arr.length) return 0;
    if (dp[i] != -1) return dp[i];
    int maxVal = 0, best = 0;
    for (int len = 1; len <= k && i + len <= arr.length; len++) {
        maxVal = Math.max(maxVal, arr[i + len - 1]);
        best = Math.max(best, maxVal * len + memo(i + len, arr, k, dp));
    }
    return dp[i] = best;
}

Complexity. Time O(n · k): n states, each looping over at most k lengths in O(1) per length thanks to the running maximum. Space O(n) for dp plus O(n) stack.

Approach 3: Tabulation

Intuition. f(i) reads only f(i + 1) .. f(i + k), all to its right. Filling dp from n - 1 down to 0 therefore always finds the needed values ready. dp[n] = 0 is the default.

Algorithm.

1. Allocate dp of size n + 1.
2. For i from n - 1 down to 0, compute dp[i] with the same inner loop.
3. Return dp[0].

static int tabulation(int[] arr, int k) {
    int n = arr.length;
    int[] dp = new int[n + 1];                              // dp[n] = 0
    for (int i = n - 1; i >= 0; i--) {
        int maxVal = 0, best = 0;
        for (int len = 1; len <= k && i + len <= n; len++) {
            maxVal = Math.max(maxVal, arr[i + len - 1]);
            best = Math.max(best, maxVal * len + dp[i + len]);
        }
        dp[i] = best;
    }
    return dp[0];
}

Dry run on Example 1, arr = [1, 15, 7, 9, 2, 5, 10], k = 3:

The winning choice at i = 0 is a piece of length 3 ([1, 15, 7], worth 45), then dp[3] = 39 comes from [9] followed by dp[4] = 30, which is [2, 5, 10].

Complexity. Time O(n · k), space O(n).

Approach 4: Space optimised, a ring of k + 1 cells

Intuition. Each dp[i] looks at most k cells ahead. Once dp[i] is computed, dp[i + k + 1] will never be read again (every later index i' < i looks only as far as i' + k < i + k + 1). So k + 1 cells are enough if we store f(i) at position i % (k + 1). At the moment we compute f(i), the ring holds f(i + 1) .. f(i + k + 1); the one we no longer need, f(i + k + 1), lives in exactly the slot i % (k + 1) that f(i) overwrites.

The base value f(n) = 0 needs no special handling: the ring starts as all zeros, and the slot for f(n) is only overwritten by f(n - k - 1), after every index that reads f(n) has finished.

Algorithm.

1. Allocate dp of size k + 1 (zeros).
2. For i from n - 1 down to 0, run the usual inner loop reading dp[(i + len) % (k + 1)], and store the result in dp[i % (k + 1)].
3. Return dp[0], which holds f(0).

static int spaceOptimized(int[] arr, int k) {
    int n = arr.length;
    int[] dp = new int[k + 1];                              // dp[i % (k + 1)] holds f(i); f(n) starts as 0
    for (int i = n - 1; i >= 0; i--) {
        int maxVal = 0, best = 0;
        for (int len = 1; len <= k && i + len <= n; len++) {
            maxVal = Math.max(maxVal, arr[i + len - 1]);
            best = Math.max(best, maxVal * len + dp[(i + len) % (k + 1)]);
        }
        dp[i % (k + 1)] = best;                             // overwrites f(i + k + 1), which is no longer needed
    }
    return dp[0];
}

Complexity. Time O(n · k), space O(k). With k ≤ n ≤ 500 the saving is small, but the ring-buffer technique is worth knowing for any DP that looks a fixed distance ahead or behind.

Edge cases and pitfalls

• k = 1. Every piece is a single element, so the answer is the plain sum: [3, 8, 2] gives 13.
• k ≥ n. With non-negative values, one piece over the whole array is optimal (max · n), since no element can ever be raised above the global maximum. [5, 1, 1] with k = 3 gives 15.
• Bounds. The inner loop must stop at i + len ≤ n; forgetting it reads past the end of the array.
• Starting maxVal at 0. This relies on arr[i] ≥ 0, which the constraints guarantee. For arrays with negative numbers, start it at Integer.MIN_VALUE.
• Overflow. Every candidate is the sum of a real partition, so it is at most the final answer, which LeetCode promises fits in an int. Without that promise (values up to 10^9, n up to 500), use long.
• Greedy does not work. Letting the largest value absorb as many neighbours as possible is tempting, but the pieces must tile the array and one choice limits the next. In [2, 9, 2, 2, 9] with k = 2, the best is 38, for example [2, 9] [2] [2, 9].
• Testing. The Java file compares all four methods on 300 seeded random arrays, and runs the polynomial methods on a 500-element array with k = 100.

Related problems

• Palindrome partitioning II has the same front-partition structure with a validity test on each piece.
• Matrix chain multiplication is the interval form of partition DP, where a split can happen anywhere in a range.
• Frog jump with K distances also reads the previous k states, giving the same O(n · k) shape.
• Rod Cutting Problem - (DP - 24) chooses piece lengths too, but with a fixed price per length.