Partition to K Equal Sum Subsets

MEDIUM✓ CHIP-TIMEDLC #698 — FULL STATEMENT ↗

The drill: A list of positive numbers and a count k — decide whether every number, used exactly once, can be sorted into k groups that all sum to the same value.

THE BRIEFING — THE FULL DRILL, IN MY OWN WORDS

A list of positive numbers and a count k arrive together. The task is to decide whether every number, used exactly once, can be sorted into k groups that all add up to the same total.

Every number must land in exactly one of the k groups — nothing left over, nothing split — and the groups don't need to be the same size, only the same sum.

Only a yes-or-no answer is needed: whether some valid grouping exists, not which grouping it is.

EX 01

nums = [7] · k = 1

→ true

K = 1, THE WHOLE ARRAY IS THE ONE SUBSET

EX 02

nums = [3, 3, 3] · k = 3

→ true

N EQUALS K, EVERY NUMBER EQUAL — ONE PER BUCKET

EX 03

nums = [3, 3, 4] · k = 3

→ false

N EQUALS K BUT THE NUMBERS AREN'T EQUAL

numbers are positive, typically a few dozen at most
k is a positive integer no larger than the count of numbers
every number is used exactly once across the k groups
groups must share the same sum but can differ in size

THE HINTS — TAKE ONLY WHAT YOU NEED

HINT 1 THE NUDGE

The same cheap filter applies here as with any equal-split puzzle: the total must divide evenly by k before any arrangement stands a chance. What's each group's target sum once that passes?

HINT 2 THE STRUCTURE

Each number drops into exactly one of k running bucket totals, and every bucket must land on the target exactly — it's bin-packing with k bins instead of a fixed four.

HINT 3 ONE STEP FROM THE ANSWER

The set of numbers already placed determines everything that matters — cache, per bitmask of used numbers, how far the current bucket has filled, mod the target. Two orders that used the same numbers always reach the same partial fill, so solve each bitmask once.

COACH'S BOARD — THE PATTERN, STEP BY STEP

THE K-BUCKET CACHEPATTERN · BITMASK DP OVER SUBSETSnums = [6, 3, 3, 2, 2, 2] · k = 3 — target 6

STEP 1

Six numbers: 6, 3, 3, 2, 2, 2. Sum 18 over k = 3 buckets — target 6 per bucket.

STEP 1 / 8 · ← → WORK TOO

THE SPLITS — TWO PACES, TWO LANGUAGES

grind/partition-to-k-equal-sum-subsets.py — RACE PACE

LANG ▸

PACE ▸

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        total = sum(nums)
        if total % k != 0:
            return False
        target = total // k
        n = len(nums)
        if any(x > target for x in nums):
            return False

        full = 1 << n
        # progress[mask] = how far the current bucket has filled, mod `target`,
        # after placing exactly the numbers in `mask`. -1 = unreachable.
        progress = [-1] * full
        progress[0] = 0
        for mask in range(full):
            if progress[mask] == -1:
                continue
            for i in range(n):
                if mask & (1 << i):
                    continue
                nxt = mask | (1 << i)
                if progress[nxt] != -1:
                    continue
                if progress[mask] + nums[i] <= target:
                    progress[nxt] = (progress[mask] + nums[i]) % target

        return progress[full - 1] == 0

TIME O(N · 2^N)SPACE O(2^N)PYTHON · RACE PACE · 28 LN

class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        total = sum(nums)
        if total % k != 0:
            return False
        target = total // k
        n = len(nums)
        buckets = [0] * k

        def backtrack(i: int) -> bool:
            if i == n:
                return all(b == target for b in buckets)
            for j in range(k):
                if buckets[j] + nums[i] <= target:
                    buckets[j] += nums[i]
                    if backtrack(i + 1):
                        return True
                    buckets[j] -= nums[i]
            return False

        return backtrack(0)

TIME O(K^N)SPACE O(N)PYTHON · TRAINING PACE · 21 LN

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;
        int n = nums.length;
        for (int x : nums) {
            if (x > target) {
                return false;
            }
        }

        int full = 1 << n;
        // progress[mask] = how far the current bucket has filled, mod `target`,
        // after placing exactly the numbers in `mask`. -1 = unreachable.
        int[] progress = new int[full];
        Arrays.fill(progress, -1);
        progress[0] = 0;
        for (int mask = 0; mask < full; mask++) {
            if (progress[mask] == -1) {
                continue;
            }
            for (int i = 0; i < n; i++) {
                if ((mask & (1 << i)) != 0) {
                    continue;
                }
                int next = mask | (1 << i);
                if (progress[next] != -1) {
                    continue;
                }
                if (progress[mask] + nums[i] <= target) {
                    progress[next] = (progress[mask] + nums[i]) % target;
                }
            }
        }

        return progress[full - 1] == 0;
    }
}

TIME O(N · 2^N)SPACE O(2^N)JAVA · RACE PACE · 44 LN

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;
        int[] buckets = new int[k];
        return backtrack(nums, 0, buckets, target);
    }

    private boolean backtrack(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 j = 0; j < buckets.length; j++) {
            if (buckets[j] + nums[i] <= target) {
                buckets[j] += nums[i];
                if (backtrack(nums, i + 1, buckets, target)) {
                    return true;
                }
                buckets[j] -= nums[i];
            }
        }
        return false;
    }
}

TIME O(K^N)SPACE O(N)JAVA · TRAINING PACE · 35 LN

✓ CHIP-TIMED — ALL 4 SOLUTIONS RAN GREEN AGAINST SELF-AUTHORED CASES IN CI · JDK 21 · CPYTHON 3.12 · NOTHING PUBLISHES RED