MediumLeetCode 1043Partition DP — choose the last segment

Partition Array for Maximum Sum (LeetCode 1043)

Problem 4 of 10 in Advanced DP — Practice Problems

The problem

You get an integer array arr and an integer k. Cut the array into contiguous blocks, each of length at most k. Then every element in a block is replaced by that block’s maximum. Return the largest possible sum of the resulting array.

Input:  arr = [1, 15, 7, 9, 2, 5, 10], k = 3
Output: 84

Cut as [1, 15, 7] [9] [2, 5, 10]. That becomes [15, 15, 15, 9, 10, 10, 10], which sums to 84.

How to think about it

Start with brute force. Between each pair of neighbours you either cut or don’t, subject to blocks of length ≤ k. That is up to 2ⁿ⁻¹ ways to cut, far too many for n = 500.

The key insight: fix the last block. Look at the first i elements. Whatever the best cutting is, its last block ends at i and has some length j between 1 and k. Everything before that block is just the same problem on the first i - j elements. So:

dp[i] = max over j in 1..k of dp[i - j] + (max of the last j elements) × j

Get the block max for free. As j grows by one, the block gains one element on its left, arr[i - j]. Keep a running cur_max while scanning j upwards, so each candidate costs O(1).

Which pattern is this? Partition DP, the “choose the last segment” template. The tell is “split an array into contiguous pieces, each piece scored on its own”, as in Word Break or Palindrome Partitioning.

Approach

  1. Create dp of size n + 1, with dp[0] = 0 (empty prefix).
  2. For each prefix length i from 1 to n:
    1. Reset cur_max = 0.
    2. For j from 1 to min(k, i), extend the last block left by one: cur_max = max(cur_max, arr[i - j]).
    3. Update dp[i] = max(dp[i], dp[i - j] + cur_max * j).
  3. Return dp[n].

Solution

def max_sum_after_partitioning(arr, k):
    n = len(arr); dp = [0] * (n + 1)
    for i in range(1, n + 1):
        cur_max = 0
        for j in range(1, min(k, i) + 1):          # last block = arr[i-j .. i-1]
            cur_max = max(cur_max, arr[i - j])      # grow it leftwards
            dp[i] = max(dp[i], dp[i - j] + cur_max * j)
    return dp[n]
# max_sum_after_partitioning([1,15,7,9,2,5,10], 3) -> 84

Dry run

On [1, 15, 7, 9, 2, 5, 10], k = 3. “Best last block” is the (length, max) that produced dp[i].

i arr[i-1] Best last block dp[i]
1 1 length 1, max 1 1
2 15 length 2, max 15 30
3 7 length 3, max 15 45
4 9 length 1, max 9 54
5 2 length 2, max 9 63
6 5 length 3, max 9 72
7 10 length 3, max 10 84

At i = 7, the last block [2, 5, 10] gives dp[4] + 10 × 3 = 54 + 30 = 84.

Complexity

  • Time: O(n·k). For each of n prefixes, the inner loop runs at most k times with O(1) work.
  • Space: O(n). For the dp array. Only the last k entries are read, so a ring buffer of size k + 1 would also work.

Edge cases and pitfalls

  • k = 1: no element can change, so the answer is sum(arr).
  • k ≥ n: one block is allowed, so the answer can be max(arr) × n.
  • Recomputing the block max with max(arr[i-j:i]) makes it O(n·k²). Keep the running max.
  • Off-by-one on indices: dp[i] covers arr[0..i-1], so the element added at step j is arr[i - j], not arr[i - j + 1].
  • Greedy “give the biggest number the widest block” fails; neighbouring blocks compete for space.