1043. Partition Array for Maximum Sum - Explanation

Description

Split arr into contiguous, nonempty groups, each containing at most k elements. Replace every value in each group with the largest original value in that group. Return the greatest total sum obtainable after these replacements.

The inputs are chosen so that the maximum sum fits in a signed 32-bit integer.

Example 1:

Input: arr = [2,6,1,4]
k = 2

Output: 20

Explanation: Use groups [2,6] and [1,4]. They become [6,6] and [4,4], giving 20. No other allowed partition produces a larger sum.

Example 2:

Input: arr = [3,1,5]
k = 1

Output: 9

Explanation: With k = 1, every group has one element. No value changes, so the sum is 3 + 1 + 5 = 9.

Example 3:

Input: arr = [0,0]
k = 2

Output: 0

Explanation: All original values are zero. Every possible group still has maximum zero, so the final sum is zero.

Constraints:

1 <= arr.length <= 500
0 <= arr[i] <= 1000000000
1 <= k <= arr.length
The maximum obtainable sum is at most 2147483647.

Topics

Array Dynamic Programming

Company Tags

Please upgrade to NeetCode Pro to view company tags.

Prerequisites

Before attempting this problem, you should be comfortable with:

Dynamic Programming (Recursion with Memoization) - Core technique for solving this optimization problem efficiently
Array Partitioning - Understanding how to split arrays into contiguous subarrays
Subarray Maximum Tracking - Needed to find the maximum element within sliding windows
Space Optimization in DP - Useful for reducing memory usage with circular buffer technique

1. Recursion

Intuition

We need to partition the array into subarrays of length at most k, where each element in a subarray becomes the maximum value of that subarray. The goal is to maximize the total sum after this transformation.

At each position, we have a choice: end the current subarray at any of the next k positions. For each choice, we calculate the contribution (maximum element times subarray length) and recursively solve the remaining array. We try all valid partition lengths and take the maximum result.

Algorithm

Define a recursive function starting at index i.
Base case: if i reaches the end of the array, return 0.
For each possible subarray ending at positions i to min(i + k - 1, n - 1):
- Track the maximum element seen so far in this subarray.
- Calculate the sum contribution as max element times window size.
- Recursively solve for the rest of the array starting at j + 1.
- Keep track of the best total sum.
Return the maximum sum found.

class Solution:
    def maxSumAfterPartitioning(self, arr: List[int], k: int) -> int:
        def dfs(i):
            if i >= len(arr):
                return 0

            cur_max = 0
            res = 0
            for j in range(i, min(len(arr), i + k)):
                cur_max = max(cur_max, arr[j])
                window_size = j - i + 1
                res = max(res, dfs(j + 1) + cur_max * window_size)

            return res

        return dfs(0)

public class Solution {
    public int maxSumAfterPartitioning(int[] arr, int k) {
        return dfs(0, arr, k);
    }

    private int dfs(int i, int[] arr, int k) {
        if (i >= arr.length) {
            return 0;
        }

        int cur_max = 0;
        int res = 0;
        for (int j = i; j < Math.min(arr.length, i + k); j++) {
            cur_max = Math.max(cur_max, arr[j]);
            int window_size = j - i + 1;
            res = Math.max(res, dfs(j + 1, arr, k) + cur_max * window_size);
        }

        return res;
    }
}

class Solution {
public:
    int maxSumAfterPartitioning(vector<int>& arr, int k) {
        return dfs(0, arr, k);
    }

private:
    int dfs(int i, vector<int>& arr, int k) {
        if (i >= arr.size()) {
            return 0;
        }

        int cur_max = 0, res = 0;
        for (int j = i; j < min((int)arr.size(), i + k); j++) {
            cur_max = max(cur_max, arr[j]);
            int window_size = j - i + 1;
            res = max(res, dfs(j + 1, arr, k) + cur_max * window_size);
        }

        return res;
    }
};

class Solution {
    /**
     * @param {number[]} arr
     * @param {number} k
     * @return {number}
     */
    maxSumAfterPartitioning(arr, k) {
        const dfs = (i) => {
            if (i >= arr.length) {
                return 0;
            }

            let cur_max = 0,
                res = 0;
            for (let j = i; j < Math.min(arr.length, i + k); j++) {
                cur_max = Math.max(cur_max, arr[j]);
                let window_size = j - i + 1;
                res = Math.max(res, dfs(j + 1) + cur_max * window_size);
            }

            return res;
        };

        return dfs(0);
    }
}

public class Solution {
    public int MaxSumAfterPartitioning(int[] arr, int k) {
        return Dfs(0, arr, k);
    }

    private int Dfs(int i, int[] arr, int k) {
        if (i >= arr.Length) {
            return 0;
        }

        int curMax = 0, res = 0;
        for (int j = i; j < Math.Min(arr.Length, i + k); j++) {
            curMax = Math.Max(curMax, arr[j]);
            int windowSize = j - i + 1;
            res = Math.Max(res, Dfs(j + 1, arr, k) + curMax * windowSize);
        }

        return res;
    }
}

func maxSumAfterPartitioning(arr []int, k int) int {
    var dfs func(i int) int
    dfs = func(i int) int {
        if i >= len(arr) {
            return 0
        }

        curMax, res := 0, 0
        for j := i; j < min(len(arr), i+k); j++ {
            curMax = max(curMax, arr[j])
            windowSize := j - i + 1
            res = max(res, dfs(j+1)+curMax*windowSize)
        }

        return res
    }

    return dfs(0)
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

class Solution {
    fun maxSumAfterPartitioning(arr: IntArray, k: Int): Int {
        fun dfs(i: Int): Int {
            if (i >= arr.size) {
                return 0
            }

            var curMax = 0
            var res = 0
            for (j in i until minOf(arr.size, i + k)) {
                curMax = maxOf(curMax, arr[j])
                val windowSize = j - i + 1
                res = maxOf(res, dfs(j + 1) + curMax * windowSize)
            }

            return res
        }

        return dfs(0)
    }
}

class Solution {
    func maxSumAfterPartitioning(_ arr: [Int], _ k: Int) -> Int {
        func dfs(_ i: Int) -> Int {
            if i >= arr.count {
                return 0
            }

            var curMax = 0
            var res = 0
            for j in i..<min(arr.count, i + k) {
                curMax = max(curMax, arr[j])
                let windowSize = j - i + 1
                res = max(res, dfs(j + 1) + curMax * windowSize)
            }

            return res
        }

        return dfs(0)
    }
}

impl Solution {
    pub fn max_sum_after_partitioning(arr: Vec<i32>, k: i32) -> i32 {
        fn dfs(i: usize, arr: &[i32], k: usize) -> i32 {
            if i >= arr.len() {
                return 0;
            }
            let mut cur_max = 0;
            let mut res = 0;
            for j in i..arr.len().min(i + k) {
                cur_max = cur_max.max(arr[j]);
                let window_size = (j - i + 1) as i32;
                res = res.max(dfs(j + 1, arr, k) + cur_max * window_size);
            }
            res
        }

        dfs(0, &arr, k as usize)
    }
}

Time & Space Complexity

Time complexity: $O(k ^ n)$
Space complexity: $O(n)$ for recursion stack.

Where $k$ is the maximum length of the subarray and $n$ is the size of the array $arr$ .

2. Dynamic Programming (Top-Down)

Intuition

The recursive solution has overlapping subproblems. When we compute the maximum sum starting from index i, we might need this value multiple times from different partition choices. By caching results, we avoid redundant calculations.

This is the memoized version of the recursive approach. We store computed results in a cache and return them directly when we encounter the same subproblem again.

Algorithm

Create a cache dictionary with base case: cache[n] = 0.
Define a recursive function with memoization starting at index i.
If i is in the cache, return the cached value immediately.
Try all possible subarray lengths from 1 to k:
- Track the maximum element in the current window.
- Calculate contribution and add the recursive result for the remaining array.
- Take the maximum across all choices.
Store the result in the cache and return it.

class Solution:
    def maxSumAfterPartitioning(self, arr: List[int], k: int) -> int:
        cache = { len(arr) : 0 }

        def dfs(i):
            if i in cache:
                return cache[i]

            cur_max = 0
            res = 0
            for j in range(i, min(len(arr), i + k)):
                cur_max = max(cur_max, arr[j])
                window_size = j - i + 1
                res = max(res, dfs(j + 1) + cur_max * window_size)

            cache[i] = res
            return res

        return dfs(0)

public class Solution {
    public int maxSumAfterPartitioning(int[] arr, int k) {
        int[] cache = new int[arr.length + 1];
        Arrays.fill(cache, -1);
        cache[arr.length] = 0;
        return dfs(0, arr, k, cache);
    }

    private int dfs(int i, int[] arr, int k, int[] cache) {
        if (cache[i] != -1) {
            return cache[i];
        }

        int cur_max = 0, res = 0;
        for (int j = i; j < Math.min(arr.length, i + k); j++) {
            cur_max = Math.max(cur_max, arr[j]);
            int window_size = j - i + 1;
            res = Math.max(res, dfs(j + 1, arr, k, cache) + cur_max * window_size);
        }

        cache[i] = res;
        return res;
    }
}

class Solution {
public:
    int maxSumAfterPartitioning(vector<int>& arr, int k) {
        vector<int> cache(arr.size() + 1, -1);
        cache[arr.size()] = 0;
        return dfs(0, arr, k, cache);
    }

private:
    int dfs(int i, vector<int>& arr, int k, vector<int>& cache) {
        if (cache[i] != -1) {
            return cache[i];
        }

        int cur_max = 0, res = 0;
        for (int j = i; j < min((int)arr.size(), i + k); j++) {
            cur_max = max(cur_max, arr[j]);
            int window_size = j - i + 1;
            res = max(res, dfs(j + 1, arr, k, cache) + cur_max * window_size);
        }

        return cache[i] = res;
    }
};

class Solution {
    /**
     * @param {number[]} arr
     * @param {number} k
     * @return {number}
     */
    maxSumAfterPartitioning(arr, k) {
        const cache = new Array(arr.length + 1).fill(-1);
        cache[arr.length] = 0;

        const dfs = (i) => {
            if (cache[i] !== -1) {
                return cache[i];
            }

            let cur_max = 0,
                res = 0;
            for (let j = i; j < Math.min(arr.length, i + k); j++) {
                cur_max = Math.max(cur_max, arr[j]);
                let window_size = j - i + 1;
                res = Math.max(res, dfs(j + 1) + cur_max * window_size);
            }

            return (cache[i] = res);
        };

        return dfs(0);
    }
}

public class Solution {
    public int MaxSumAfterPartitioning(int[] arr, int k) {
        int[] cache = new int[arr.Length + 1];
        Array.Fill(cache, -1);
        cache[arr.Length] = 0;
        return Dfs(0, arr, k, cache);
    }

    private int Dfs(int i, int[] arr, int k, int[] cache) {
        if (cache[i] != -1) {
            return cache[i];
        }

        int curMax = 0, res = 0;
        for (int j = i; j < Math.Min(arr.Length, i + k); j++) {
            curMax = Math.Max(curMax, arr[j]);
            int windowSize = j - i + 1;
            res = Math.Max(res, Dfs(j + 1, arr, k, cache) + curMax * windowSize);
        }

        cache[i] = res;
        return res;
    }
}

func maxSumAfterPartitioning(arr []int, k int) int {
    cache := make([]int, len(arr)+1)
    for i := range cache {
        cache[i] = -1
    }
    cache[len(arr)] = 0

    var dfs func(i int) int
    dfs = func(i int) int {
        if cache[i] != -1 {
            return cache[i]
        }

        curMax, res := 0, 0
        for j := i; j < min(len(arr), i+k); j++ {
            curMax = max(curMax, arr[j])
            windowSize := j - i + 1
            res = max(res, dfs(j+1)+curMax*windowSize)
        }

        cache[i] = res
        return res
    }

    return dfs(0)
}

func min(a, b int) int {
    if a < b {
        return a
    }
    return b
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

class Solution {
    fun maxSumAfterPartitioning(arr: IntArray, k: Int): Int {
        val cache = IntArray(arr.size + 1) { -1 }
        cache[arr.size] = 0

        fun dfs(i: Int): Int {
            if (cache[i] != -1) {
                return cache[i]
            }

            var curMax = 0
            var res = 0
            for (j in i until minOf(arr.size, i + k)) {
                curMax = maxOf(curMax, arr[j])
                val windowSize = j - i + 1
                res = maxOf(res, dfs(j + 1) + curMax * windowSize)
            }

            cache[i] = res
            return res
        }

        return dfs(0)
    }
}

class Solution {
    func maxSumAfterPartitioning(_ arr: [Int], _ k: Int) -> Int {
        var cache = [Int](repeating: -1, count: arr.count + 1)
        cache[arr.count] = 0

        func dfs(_ i: Int) -> Int {
            if cache[i] != -1 {
                return cache[i]
            }

            var curMax = 0
            var res = 0
            for j in i..<min(arr.count, i + k) {
                curMax = max(curMax, arr[j])
                let windowSize = j - i + 1
                res = max(res, dfs(j + 1) + curMax * windowSize)
            }

            cache[i] = res
            return res
        }

        return dfs(0)
    }
}

impl Solution {
    pub fn max_sum_after_partitioning(arr: Vec<i32>, k: i32) -> i32 {
        let n = arr.len();
        let k = k as usize;
        let mut cache = vec![-1; n + 1];
        cache[n] = 0;

        fn dfs(i: usize, arr: &[i32], k: usize, cache: &mut [i32]) -> i32 {
            if cache[i] != -1 {
                return cache[i];
            }
            let mut cur_max = 0;
            let mut res = 0;
            for j in i..arr.len().min(i + k) {
                cur_max = cur_max.max(arr[j]);
                let window_size = (j - i + 1) as i32;
                res = res.max(dfs(j + 1, arr, k, cache) + cur_max * window_size);
            }
            cache[i] = res;
            res
        }

        dfs(0, &arr, k, &mut cache)
    }
}

Time & Space Complexity

Time complexity: $O(n * k)$
Space complexity: $O(n)$

Where $k$ is the maximum length of the subarray and $n$ is the size of the array $arr$ .

3. Dynamic Programming (Bottom-Up)

Intuition

We can flip the top-down approach to bottom-up by filling the DP table from right to left. dp[i] represents the maximum sum achievable for the subarray starting at index i.

Starting from the last element and working backwards, we build up solutions to larger subproblems using already-computed smaller ones. This eliminates recursion overhead and makes the memory access pattern more predictable.

Algorithm

Create a DP array of size n + 1 initialized to 0 (base case is dp[n] = 0).
Iterate from i = n - 1 down to 0:
- Track the maximum element for the current window starting at i.
- For each valid window ending at j (up to i + k - 1):
  - Update the maximum element.
  - Calculate the sum as max element times window size plus dp[j + 1].
  - Update dp[i] with the maximum value found.
Return dp[0] as the answer.

class Solution:
    def maxSumAfterPartitioning(self, arr: List[int], k: int) -> int:
        n = len(arr)
        dp = [0] * (n + 1)

        for i in range(n - 1, -1, -1):
            cur_max = 0
            for j in range(i, min(n, i + k)):
                cur_max = max(cur_max, arr[j])
                window_size = j - i + 1
                dp[i] = max(dp[i], dp[j + 1] + cur_max * window_size)

        return dp[0]

public class Solution {
    public int maxSumAfterPartitioning(int[] arr, int k) {
        int n = arr.length;
        int[] dp = new int[n + 1];

        for (int i = n - 1; i >= 0; i--) {
            int cur_max = 0;
            for (int j = i; j < Math.min(n, i + k); j++) {
                cur_max = Math.max(cur_max, arr[j]);
                int window_size = j - i + 1;
                dp[i] = Math.max(dp[i], dp[j + 1] + cur_max * window_size);
            }
        }

        return dp[0];
    }
}

class Solution {
public:
    int maxSumAfterPartitioning(vector<int>& arr, int k) {
        int n = arr.size();
        vector<int> dp(n + 1, 0);

        for (int i = n - 1; i >= 0; i--) {
            int cur_max = 0;
            for (int j = i; j < min(n, i + k); j++) {
                cur_max = max(cur_max, arr[j]);
                int window_size = j - i + 1;
                dp[i] = max(dp[i], dp[j + 1] + cur_max * window_size);
            }
        }

        return dp[0];
    }
};

class Solution {
    /**
     * @param {number[]} arr
     * @param {number} k
     * @return {number}
     */
    maxSumAfterPartitioning(arr, k) {
        let n = arr.length;
        let dp = new Array(n + 1).fill(0);

        for (let i = n - 1; i >= 0; i--) {
            let cur_max = 0;
            for (let j = i; j < Math.min(n, i + k); j++) {
                cur_max = Math.max(cur_max, arr[j]);
                let window_size = j - i + 1;
                dp[i] = Math.max(dp[i], dp[j + 1] + cur_max * window_size);
            }
        }

        return dp[0];
    }
}

public class Solution {
    public int MaxSumAfterPartitioning(int[] arr, int k) {
        int n = arr.Length;
        int[] dp = new int[n + 1];

        for (int i = n - 1; i >= 0; i--) {
            int curMax = 0;
            for (int j = i; j < Math.Min(n, i + k); j++) {
                curMax = Math.Max(curMax, arr[j]);
                int windowSize = j - i + 1;
                dp[i] = Math.Max(dp[i], dp[j + 1] + curMax * windowSize);
            }
        }

        return dp[0];
    }
}

class Solution {
    fun maxSumAfterPartitioning(arr: IntArray, k: Int): Int {
        val n = arr.size
        val dp = IntArray(n + 1)

        for (i in n - 1 downTo 0) {
            var curMax = 0
            for (j in i until minOf(n, i + k)) {
                curMax = maxOf(curMax, arr[j])
                val windowSize = j - i + 1
                dp[i] = maxOf(dp[i], dp[j + 1] + curMax * windowSize)
            }
        }

        return dp[0]
    }
}

class Solution {
    func maxSumAfterPartitioning(_ arr: [Int], _ k: Int) -> Int {
        let n = arr.count
        var dp = [Int](repeating: 0, count: n + 1)

        for i in stride(from: n - 1, through: 0, by: -1) {
            var curMax = 0
            for j in i..<min(n, i + k) {
                curMax = max(curMax, arr[j])
                let windowSize = j - i + 1
                dp[i] = max(dp[i], dp[j + 1] + curMax * windowSize)
            }
        }

        return dp[0]
    }
}

impl Solution {
    pub fn max_sum_after_partitioning(arr: Vec<i32>, k: i32) -> i32 {
        let n = arr.len();
        let k = k as usize;
        let mut dp = vec![0; n + 1];

        for i in (0..n).rev() {
            let mut cur_max = 0;
            for j in i..n.min(i + k) {
                cur_max = cur_max.max(arr[j]);
                let window_size = (j - i + 1) as i32;
                dp[i] = dp[i].max(dp[j + 1] + cur_max * window_size);
            }
        }

        dp[0]
    }
}

Time & Space Complexity

Time complexity: $O(n * k)$
Space complexity: $O(n)$

Where $k$ is the maximum length of the subarray and $n$ is the size of the array $arr$ .

4. Dynamic Programming (Space Optimized)

Intuition

Looking at the bottom-up solution, we notice that dp[i] only depends on dp[i+1] through dp[i+k]. We never need values more than k positions ahead. This means we can reduce our space from O(n) to O(k) using a circular buffer.

We use modulo arithmetic to wrap around and reuse array positions. This technique is common when the recurrence relation has a bounded look-ahead.

Algorithm

Create a DP array of size k (using circular indexing).
Initialize dp[0] = arr[0] as the base case.
Iterate from i = 1 to n - 1:
- For each position, look backwards at windows of size 1 to k.
- Track the maximum element in each window.
- Compute the sum using circular array indexing: dp[(j-1) % k] for the previous subproblem.
- Store the best result at dp[i % k].
Return dp[(n-1) % k] as the answer.

class Solution:
    def maxSumAfterPartitioning(self, arr: List[int], k: int) -> int:
        dp = [0] * k
        dp[0] = arr[0]

        for i in range(1, len(arr)):
            cur_max = 0
            max_at_i = 0
            for j in range(i, i - k, -1):
                if j < 0:
                    break
                cur_max = max(cur_max, arr[j])
                window_size = i - j + 1
                cur_sum = cur_max * window_size
                sub_sum = dp[(j - 1) % k] if j > 0 else 0
                max_at_i = max(max_at_i, cur_sum + sub_sum)

            dp[i % k] = max_at_i

        return dp[(len(arr) - 1) % k]

public class Solution {
    public int maxSumAfterPartitioning(int[] arr, int k) {
        int n = arr.length;
        int[] dp = new int[k];
        dp[0] = arr[0];

        for (int i = 1; i < n; i++) {
            int cur_max = 0, max_at_i = 0;
            for (int j = i; j > i - k; j--) {
                if (j < 0) break;
                cur_max = Math.max(cur_max, arr[j]);
                int window_size = i - j + 1;
                int cur_sum = cur_max * window_size;
                int sub_sum = (j > 0) ? dp[(j - 1) % k] : 0;
                max_at_i = Math.max(max_at_i, cur_sum + sub_sum);
            }
            dp[i % k] = max_at_i;
        }

        return dp[(n - 1) % k];
    }
}

class Solution {
public:
    int maxSumAfterPartitioning(vector<int>& arr, int k) {
        int n = arr.size();
        vector<int> dp(k);
        dp[0] = arr[0];

        for (int i = 1; i < n; i++) {
            int cur_max = 0, max_at_i = 0;
            for (int j = i; j > i - k; j--) {
                if (j < 0) break;
                cur_max = max(cur_max, arr[j]);
                int window_size = i - j + 1;
                int cur_sum = cur_max * window_size;
                int sub_sum = (j > 0) ? dp[(j - 1) % k] : 0;
                max_at_i = max(max_at_i, cur_sum + sub_sum);
            }
            dp[i % k] = max_at_i;
        }

        return dp[(n - 1) % k];
    }
};

class Solution {
    /**
     * @param {number[]} arr
     * @param {number} k
     * @return {number}
     */
    maxSumAfterPartitioning(arr, k) {
        const n = arr.length;
        const dp = new Array(k).fill(0);
        dp[0] = arr[0];

        for (let i = 1; i < n; i++) {
            let cur_max = 0,
                max_at_i = 0;
            for (let j = i; j > i - k; j--) {
                if (j < 0) break;
                cur_max = Math.max(cur_max, arr[j]);
                let window_size = i - j + 1;
                let cur_sum = cur_max * window_size;
                let sub_sum = j > 0 ? dp[(j - 1) % k] : 0;
                max_at_i = Math.max(max_at_i, cur_sum + sub_sum);
            }
            dp[i % k] = max_at_i;
        }

        return dp[(n - 1) % k];
    }
}

public class Solution {
    public int MaxSumAfterPartitioning(int[] arr, int k) {
        int n = arr.Length;
        int[] dp = new int[k];
        dp[0] = arr[0];

        for (int i = 1; i < n; i++) {
            int curMax = 0, maxAtI = 0;
            for (int j = i; j > i - k; j--) {
                if (j < 0) break;
                curMax = Math.Max(curMax, arr[j]);
                int windowSize = i - j + 1;
                int curSum = curMax * windowSize;
                int subSum = j > 0 ? dp[(j - 1) % k] : 0;
                maxAtI = Math.Max(maxAtI, curSum + subSum);
            }
            dp[i % k] = maxAtI;
        }

        return dp[(n - 1) % k];
    }
}

func maxSumAfterPartitioning(arr []int, k int) int {
    n := len(arr)
    dp := make([]int, k)
    dp[0] = arr[0]

    for i := 1; i < n; i++ {
        curMax, maxAtI := 0, 0
        for j := i; j > i-k; j-- {
            if j < 0 {
                break
            }
            curMax = max(curMax, arr[j])
            windowSize := i - j + 1
            curSum := curMax * windowSize
            subSum := 0
            if j > 0 {
                subSum = dp[(j-1)%k]
            }
            maxAtI = max(maxAtI, curSum+subSum)
        }
        dp[i%k] = maxAtI
    }

    return dp[(n-1)%k]
}

func max(a, b int) int {
    if a > b {
        return a
    }
    return b
}

class Solution {
    fun maxSumAfterPartitioning(arr: IntArray, k: Int): Int {
        val n = arr.size
        val dp = IntArray(k)
        dp[0] = arr[0]

        for (i in 1 until n) {
            var curMax = 0
            var maxAtI = 0
            var j = i
            while (j > i - k) {
                if (j < 0) break
                curMax = maxOf(curMax, arr[j])
                val windowSize = i - j + 1
                val curSum = curMax * windowSize
                val subSum = if (j > 0) dp[(j - 1) % k] else 0
                maxAtI = maxOf(maxAtI, curSum + subSum)
                j--
            }
            dp[i % k] = maxAtI
        }

        return dp[(n - 1) % k]
    }
}

class Solution {
    func maxSumAfterPartitioning(_ arr: [Int], _ k: Int) -> Int {
        let n = arr.count
        var dp = [Int](repeating: 0, count: k)
        dp[0] = arr[0]

        for i in 1..<n {
            var curMax = 0
            var maxAtI = 0
            var j = i
            while j > i - k {
                if j < 0 { break }
                curMax = max(curMax, arr[j])
                let windowSize = i - j + 1
                let curSum = curMax * windowSize
                let subSum = j > 0 ? dp[(j - 1) % k] : 0
                maxAtI = max(maxAtI, curSum + subSum)
                j -= 1
            }
            dp[i % k] = maxAtI
        }

        return dp[(n - 1) % k]
    }
}

impl Solution {
    pub fn max_sum_after_partitioning(arr: Vec<i32>, k: i32) -> i32 {
        let n = arr.len();
        let k = k as usize;
        let mut dp = vec![0i32; k];
        dp[0] = arr[0];

        for i in 1..n {
            let mut cur_max = 0;
            let mut max_at_i = 0;
            let mut j = i as i32;
            while j > i as i32 - k as i32 {
                if j < 0 {
                    break;
                }
                let ju = j as usize;
                cur_max = cur_max.max(arr[ju]);
                let window_size = (i - ju + 1) as i32;
                let cur_sum = cur_max * window_size;
                let sub_sum = if ju > 0 { dp[(ju - 1) % k] } else { 0 };
                max_at_i = max_at_i.max(cur_sum + sub_sum);
                j -= 1;
            }
            dp[i % k] = max_at_i;
        }

        dp[(n - 1) % k]
    }
}

Time & Space Complexity

Time complexity: $O(n * k)$
Space complexity: $O(k)$

Where $k$ is the maximum length of the subarray and $n$ is the size of the array $arr$ .

Common Pitfalls

Misunderstanding the Transformation

Each element in a partition becomes the maximum value of that partition, not stays as its original value. The sum contribution is max_element * partition_length, not the sum of original elements. This transformation is the key insight of the problem.

Incorrect Window Size Calculation

When iterating through possible partition endpoints, the window size is j - i + 1, not j - i. Off-by-one errors here lead to underestimating the contribution of each partition and returning a smaller sum than optimal.

Forgetting to Track Running Maximum

For each partition starting at index i, you must track the maximum element seen so far as you extend the window. Recalculating the maximum from scratch for each window length turns an O(nk) solution into O(nk^2).