Partition to K Equal Sum Subsets

NeetCode: Partition to K Equal Sum Subsets · LeetCode: 698. Partition to K Equal Sum Subsets · Difficulty: Medium

Given an integer array nums and an integer k, return true if it is possible to divide the array into k non-empty subsets whose sums are all equal.

Example

Text

Input:  nums = [4,3,2,3,5,2,1], k = 4   Output: true   // (5), (1,4), (2,3), (2,3)
Input:  nums = [1,2,3,4], k = 3         Output: false

Both approaches first reject two cases: the total must be divisible by k, and no number may exceed the target total / k.

Brute force: try each number in each bucket

Keep k bucket sums. For each number in turn, try putting it in every bucket whose sum stays within the target. The partition is valid if all numbers are placed and every bucket ends at the target.

LeetCode's Java signature is boolean canPartitionKSubsets(int[] nums, int k). The sums are int values, which wrap silently on overflow, so very large numbers would need long.

Java

public class Main {
    static class Solution {
        public boolean canPartitionKSubsets(int[] nums, int k) {
            int total = 0;
            for (int x : nums) {
                total += x;
            }
            if (total % k != 0) {
                return false;
            }
            int target = total / k;
            for (int x : nums) {
                if (x > target) {
                    return false;
                }
            }
            return place(nums, 0, new int[k], target);
        }

        private boolean place(int[] nums, int i, int[] buckets, int target) {
            if (i == nums.length) {
                for (int b : buckets) {
                    if (b != target) {
                        return false;
                    }
                }
                return true;
            }
            for (int b = 0; b < buckets.length; b++) {
                if (buckets[b] + nums[i] <= target) {
                    buckets[b] += nums[i];
                    if (place(nums, i + 1, buckets, target)) {
                        return true;
                    }
                    buckets[b] -= nums[i];
                }
            }
            return false;
        }
    }

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

Output

true
false
false
true

Every number has up to k choices, and buckets with equal sums are interchangeable, so symmetric states are explored repeatedly. Time: O(k^n) worst case. Space: O(n + k).

Optimal: fill the buckets in order, with two prunings

The same bucket search, with three prunings:

Sort descending. Large numbers are the most constrained, so placing them first exposes dead ends early.
Skip equal buckets. If bucket b has the same current sum as an earlier bucket p < b, trying b repeats what trying p already did and failed, so skip it.
Stop when a number fits nowhere. If no bucket leads to success, return false immediately.

The equal-sum skip is safe because the attempt with p already explored every completion of the remaining numbers, and swapping the contents of two equal-sum buckets gives an equivalent state. A success through b would therefore have been a success through p. In particular, only the first empty bucket is ever tried.

Arrays.sort(int[]) sorts ascending, so the method sorts and then reads the array from the end: index n - 1 is the largest number, and the search walks i from n - 1 down to 0.

Java

import java.util.Arrays;

public class Main {
    static class Solution {
        public boolean canPartitionKSubsets(int[] nums, int k) {
            int total = 0;
            for (int x : nums) {
                total += x;
            }
            if (total % k != 0) {
                return false;
            }
            int target = total / k;
            Arrays.sort(nums); // ascending; the search reads from the largest end
            if (nums[nums.length - 1] > target) {
                return false;
            }
            return place(nums, nums.length - 1, new int[k], target);
        }

        private boolean place(int[] nums, int i, int[] sums, int target) {
            if (i < 0) {
                return true; // every number fit and total == k * target, so all buckets are full
            }
            for (int b = 0; b < sums.length; b++) {
                if (sums[b] + nums[i] > target) {
                    continue; // does not fit
                }
                boolean seen = false;
                for (int p = 0; p < b; p++) {
                    if (sums[p] == sums[b]) {
                        seen = true; // identical to an earlier bucket
                        break;
                    }
                }
                if (seen) {
                    continue;
                }
                sums[b] += nums[i];
                if (place(nums, i - 1, sums, target)) {
                    return true;
                }
                sums[b] -= nums[i];
            }
            return false;
        }
    }

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

Output

true
false
false
true

Returning true once every number is placed is correct: no bucket exceeds target, all n numbers are placed, and the sums add up to k * target, so every bucket equals target.

Time: O(k^n) in the worst case, far less in practice thanks to the prunings. Space: O(n + k). A bitmask DP over the set of used numbers gives a guaranteed O(n * 2^n) bound.

How to explain it in an interview

Baseline: give each number to one of k buckets, capped at total / k: O(k^n).
Key insight: buckets with equal sums are interchangeable, so trying more than one of them is wasted work; placing large numbers first fails faster.
Steps: (1) reject if total % k != 0 or the max exceeds target; (2) sort and process from the largest; (3) place(i) returns true when every number is placed; (4) for each bucket, skip if it overflows or equals an earlier bucket's sum; (5) add, recurse, undo.
Complexity: exponential worst case, fast in practice with the prunings; space O(n + k).
Follow-ups and pitfalls: a guaranteed bound comes from bitmask DP in O(n * 2^n). Omitting the equal-sum skip causes time-limit failures on inputs such as many 1s with a large k.