Split Array Largest Sum

NeetCode: Split Array Largest Sum · LeetCode: 410. Split Array Largest Sum · Difficulty: Hard

Given an integer array nums (non-negative values) and an integer k, split nums into k non-empty contiguous subarrays so that the largest sum among these subarrays is as small as possible. Return that minimised largest sum. LeetCode's constraints (at most 1000 elements, each at most 10^6) keep every sum at or below 10^9, so the answer fits in an int. The code still adds in long so that a sum cannot wrap silently.

Example

Text

Input:  nums = [7,2,5,10,8], k = 2    Output: 18   ([7,2,5] and [10,8])
Input:  nums = [1,2,3,4,5], k = 2     Output: 9    ([1,2,3] and [4,5])
Input:  nums = [1,4,4],     k = 3     Output: 4

Brute force: dynamic programming over split points

Let dp[j][i] be the smallest possible largest sum when the first i elements are split into j parts. The last part is nums[p..i] for some split point p, so

dp[j][i] = min over p of max(dp[j-1][p], sum(nums[p..i]))

with prefix sums giving each range sum in O(1). The base case is dp[1][i] = prefix[i]. The answer is dp[k][n]. Every state tries every split point, which gives the cost below. Long.MAX_VALUE marks a state that has no valid split yet.

Java

class Solution {
    public int splitArray(int[] nums, int k) {
        int n = nums.length;
        if (k <= 0 || k > n) {
            return -1; // impossible: parts must be non-empty
        }
        long[] prefix = new long[n + 1];
        for (int i = 0; i < n; i++) {
            prefix[i + 1] = prefix[i] + nums[i];
        }
        // dp[j][i]: first i elements into j parts.
        long[][] dp = new long[k + 1][n + 1];
        for (long[] row : dp) {
            java.util.Arrays.fill(row, Long.MAX_VALUE);
        }
        for (int i = 1; i <= n; i++) {
            dp[1][i] = prefix[i];
        }
        for (int j = 2; j <= k; j++) {
            for (int i = j; i <= n; i++) {
                // Last part is nums[p..i]; the first j-1 parts need at least j-1 elements.
                for (int p = j - 1; p < i; p++) {
                    long last = prefix[i] - prefix[p];
                    long cand = Math.max(dp[j - 1][p], last);
                    if (cand < dp[j][i]) {
                        dp[j][i] = cand;
                    }
                }
            }
        }
        return (int) dp[k][n];
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.splitArray(new int[] {7, 2, 5, 10, 8}, 2));
        System.out.println(s.splitArray(new int[] {1, 2, 3, 4, 5}, 2));
        System.out.println(s.splitArray(new int[] {1, 4, 4}, 3));
        System.out.println(s.splitArray(new int[] {5}, 1));
    }
}

Output

Time: O(k * n^2). Space: O(k * n) for the table.

Optimal: binary search on the answer with a `canSplit` check

Instead of choosing the splits, guess the answer limit (the largest allowed part sum) and ask: can the array be cut into at most k parts, each with sum <= limit?

canSplit(limit) is greedy: sweep left to right, keep adding to the current part, and start a new part only when the next element would push the sum over limit. Packing as much as possible into each part never uses more parts than any other packing.

If limit works, any larger limit works too, so binary search the smallest feasible one in [max(nums), sum(nums)]. The lower bound is the largest single element (it must fit somewhere); the upper bound is one part holding everything. "At most k parts" gives the same answer as "exactly k", because a part can always be split further (when k <= n) without raising the largest sum.

For [7,2,5,10,8], k = 2: range [10, 32]. limit = 21 gives [7,2,5], [10,8]: 2 parts, feasible, so try lower. limit = 15 gives [7,2,5], [10], [8]: 3 parts, not feasible, so try higher. The smallest feasible value is 18.

Java

class Solution {
    // Greedy: can nums be cut into at most k parts, each with sum <= limit?
    private boolean canSplit(int[] nums, int k, long limit) {
        int parts = 1;
        long current = 0;
        for (int x : nums) {
            if (current + x > limit) {
                parts++; // start a new part
                current = x;
            } else {
                current += x;
            }
        }
        return parts <= k;
    }

    public int splitArray(int[] nums, int k) {
        if (k <= 0 || k > nums.length) {
            return -1; // impossible: parts must be non-empty
        }
        long lo = 0;
        long hi = 0;
        for (int x : nums) {
            lo = Math.max(lo, x);
            hi += x;
        }
        while (lo < hi) {
            long mid = lo + (hi - lo) / 2;
            if (canSplit(nums, k, mid)) {
                hi = mid; // feasible: try smaller
            } else {
                lo = mid + 1;
            }
        }
        return (int) lo;
    }
}

public class Main {
    public static void main(String[] args) {
        Solution s = new Solution();
        System.out.println(s.splitArray(new int[] {7, 2, 5, 10, 8}, 2));
        System.out.println(s.splitArray(new int[] {1, 2, 3, 4, 5}, 2));
        System.out.println(s.splitArray(new int[] {1, 4, 4}, 3));
        System.out.println(s.splitArray(new int[] {5}, 1));
        System.out.println(s.splitArray(new int[] {0, 0, 0}, 2));
    }
}

Output

Time: O(n * log(sum)): about log2(sum) guesses, each a linear sweep. Space: O(1).

How to explain it in an interview

Baseline: DP over split points, where dp[j][i] is the best largest sum for the first i elements in j parts; it is O(k * n^2).
Key idea: checking whether the largest part can be at most limit is a greedy sweep, and feasibility is monotone in limit, so binary search the answer instead of the splits.
Steps: (1) canSplit(limit): sweep, open a new part when the sum would exceed limit, compare the part count with k; (2) set lo = max, hi = sum; (3) binary search: feasible means hi = mid, else lo = mid + 1; (4) return lo.
Complexity: O(n log(sum)) time, O(1) space.
Common questions: lo = max(nums) because some part must hold the largest element. The sum uses long so that large inputs cannot wrap around. The same approach solves Capacity To Ship Packages Within D Days.

Brute force: dynamic programming over split points

Optimal: binary search on the answer with a canSplit check

How to explain it in an interview

Optimal: binary search on the answer with a `canSplit` check