Partition of a set into K subsets with equal sum

The problem, simply

You get a list of positive numbers and a number k. Can you split all the numbers into k groups so that every group adds up to the same total? Every number must go into exactly one group.

Answer true or false.

A story to remember it

Picture it

Examples

Way 1: the simple way

Take the numbers one at a time and try putting each into every bucket that still has room.

1. Work out the target: total divided by k. If it does not divide, answer false.
2. For number i, try each bucket. If the bucket plus the number fits under the target, put it there and go to number i + 1.
3. If that fails, take the number back out and try the next bucket.
4. When every number is placed, all buckets are exactly full, so answer true.

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

    private boolean place(int i, int[] nums, int[] bucket, int target) {
        if (i == nums.length) {
            for (int b : bucket) if (b != target) return false;
            return true;
        }
        for (int b = 0; b < bucket.length; b++) {
            if (bucket[b] + nums[i] > target) continue;   // no room in this bucket
            bucket[b] += nums[i];
            if (place(i + 1, nums, bucket, target)) return true;
            bucket[b] -= nums[i];                         // take it back out
        }
        return false;
    }
}

How long it takes: each of the n numbers can go into any of the k buckets, so O(k^n) in the worst case. It also tries the same split many times, because the buckets look alike: putting a 5 in bucket 1 or in bucket 2 is really the same thing. Extra room: the buckets and the call stack, O(n) when k is at most n.

Way 2: the clever way

Fill one bucket at a time, and describe the whole situation with one thing: which numbers are already used.

Why is that enough? If you know which numbers are used, you know their total. The full buckets are full, and whatever is left over, total % target, is sitting in the bucket you are filling now. So the same set of used numbers always means the same situation, and if it failed once, it fails every time.

1. Sort the numbers from small to big.
2. Keep mask, a row of on and off switches, one per number: on means used.
3. From a mask, try each unused number. If it does not fit in the current bucket, break, because the rest are bigger.
4. Add it. If the bucket becomes exactly full, the next bucket starts at 0.
5. If every number is used, answer true. If no choice works, write "fails" for this mask in the notebook.

import java.util.*;

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);                                    // small to big
        if (nums[nums.length - 1] > target) return false;     // one number is too big alone
        Boolean[] memo = new Boolean[1 << nums.length];
        return fill(0, 0, nums, target, memo);
    }

    // mask: which numbers are used. sum: how full the current bucket is.
    private boolean fill(int mask, int sum, int[] nums, int target, Boolean[] memo) {
        if (mask == (1 << nums.length) - 1) return true;      // every number placed
        if (memo[mask] != null) return memo[mask];            // seen this situation before
        boolean ok = false;
        for (int i = 0; i < nums.length && !ok; i++) {
            if ((mask >> i & 1) == 1) continue;               // already used
            if (sum + nums[i] > target) break;                // the rest are even bigger
            int next = (sum + nums[i]) % target;              // full bucket: start the next at 0
            ok = fill(mask | (1 << i), next, nums, target, memo);
        }
        memo[mask] = ok;
        return ok;
    }
}

How long it takes: there are 2^n different masks, each one is worked out once, and each tries up to n numbers. That is O(n * 2^n). Extra room: the notebook has one entry per mask, O(2^n). This works for n up to about 20.

Watch out for

• Forgetting the % target step. When a bucket is exactly full, the next bucket must start at 0.
• Using break without sorting. The early stop only works if every later number is at least as big.
• Memo on (mask, sum) is fine but wasteful. The sum is already decided by the mask.
• Negative numbers or zeros break the "too big, stop" logic. The problem says positive; say so.

What interviewers ask next

• Can Way 1 be made faster without bitmasks? Yes. Sort from big to small so big numbers fail early, and when a number does not work in an empty bucket, do not try the other empty buckets: they are all the same. These two tricks make the bucket method fast in practice.
• What if k = 2? That is the equal halves problem. Use the subset sum table: can some numbers add up to total / 2? That takes O(n * total).
• Print the groups, not only true or false. Keep the order in which numbers were taken on the successful path, and cut it each time the running sum hits the target.
• Why is it hard in general? Even k = 2 with huge numbers is NP-complete, so exponential time in n (or time growing with the size of the numbers) is expected.

Remember this

Grown-up words: the groups are subsets, the row of switches is a bitmask, the notebook is memoization, the whole method is backtracking with memoization (a bitmask DP), and the early stops are pruning.