LeetCode 698: Partition to K Equal Sum Subsets

Problem Restatement

We are given an integer array nums and an integer k.

We need to return true if it is possible to divide nums into k non-empty subsets such that every subset has the same sum. Otherwise, return false.

Each number must be used exactly once. The official constraints include 1 <= k <= nums.length <= 16 and positive integers in nums. (leetcode.com)

Input and Output

Example function shape:

def canPartitionKSubsets(nums: list[int], k: int) -> bool:
    ...

Examples

Example 1:

nums = [4, 3, 2, 3, 5, 2, 1]
k = 4

The total sum is:

4 + 3 + 2 + 3 + 5 + 2 + 1 = 20

Each subset must have sum:

20 / 4 = 5

One valid partition is:

[5]
[4, 1]
[3, 2]
[3, 2]

Answer:

True

Example 2:

nums = [1, 2, 3, 4]
k = 3

The total sum is:

10

Since 10 is not divisible by 3, equal subset sums are impossible.

Answer:

False

First Thought: Try All Partitions

A direct solution is to try assigning every number to one of the k subsets.

For each number, there are up to k choices.

So the raw search space is roughly:

k^n

where n = len(nums).

That is large, but n <= 16, so backtracking with good pruning can pass.

Key Insight

If the array can be split into k equal-sum subsets, then every subset must sum to:

target = sum(nums) / k

So we can immediately reject when:

sum(nums) % k != 0

Also, if any number is larger than target, it cannot fit in any subset.

After that, we try to place numbers into buckets whose sums should all become target.

Sorting the numbers in descending order helps because large numbers are harder to place. If a large number cannot fit, we fail early.

Algorithm

1. Compute total = sum(nums).
2. If total % k != 0, return False.
3. Set target = total // k.
4. Sort nums in descending order.
5. If nums[0] > target, return False.
6. Create k bucket sums initialized to 0.
7. Use backtracking to place each number into one bucket:
   • If adding the number exceeds target, skip that bucket.
   • Otherwise, place it and recurse.
   • If recursion succeeds, return True.
   • Undo the placement.
   • If the bucket was empty before trying this number, break to avoid symmetric duplicate states.
8. Return whether the search succeeds.

Walkthrough

Consider:

nums = [4, 3, 2, 3, 5, 2, 1]
k = 4

Total:

20

Target per subset:

5

Sort descending:

[5, 4, 3, 3, 2, 2, 1]

Start with four empty buckets:

[0, 0, 0, 0]

Place 5:

[5, 0, 0, 0]

Place 4:

[5, 4, 0, 0]

Place 3:

[5, 4, 3, 0]

Place next 3:

[5, 4, 3, 3]

Place 2 into the bucket with 3:

[5, 4, 5, 3]

Place next 2 into the other bucket with 3:

[5, 4, 5, 5]

Place 1 into the bucket with 4:

[5, 5, 5, 5]

All buckets reach the target, so the answer is true.

Correctness

If sum(nums) is not divisible by k, then no equal-sum partition can exist, because all k subset sums would have to add up to the total. The algorithm correctly returns false in that case.

After computing target, every valid subset must sum to exactly target.

The backtracking search assigns each number to one bucket. It only permits assignments that keep each bucket sum at most target.

Therefore, every complete assignment produced by the algorithm uses every number exactly once and has no bucket exceeding target.

If all numbers are assigned, then the total sum of all buckets is k * target, and every bucket is at most target. Hence every bucket must equal target.

So any successful search path gives a valid partition.

Conversely, suppose a valid partition exists. The search tries every possible placement of each number into a bucket, except symmetric placements into equivalent empty buckets. Skipping equivalent empty buckets does not remove any distinct partition, because empty buckets have no identity. Therefore, the search can still follow a placement corresponding to the valid partition and return true.

Thus, the algorithm returns true exactly when a valid equal-sum partition exists.

Complexity

Let n = len(nums).

The pruning is important in practice:

Implementation

class Solution:
    def canPartitionKSubsets(self, nums: list[int], k: int) -> bool:
        total = sum(nums)

        if total % k != 0:
            return False

        target = total // k
        nums.sort(reverse=True)

        if nums[0] > target:
            return False

        buckets = [0] * k

        def backtrack(index: int) -> bool:
            if index == len(nums):
                return True

            value = nums[index]

            for bucket_index in range(k):
                if buckets[bucket_index] + value > target:
                    continue

                buckets[bucket_index] += value

                if backtrack(index + 1):
                    return True

                buckets[bucket_index] -= value

                if buckets[bucket_index] == 0:
                    break

            return False

        return backtrack(0)

Code Explanation

We first compute the total:

total = sum(nums)

If it cannot be split evenly:

if total % k != 0:
    return False

then the answer is impossible.

The target sum for each subset is:

target = total // k

Sorting descending improves pruning:

nums.sort(reverse=True)

If the largest number is already too large:

if nums[0] > target:
    return False

it cannot belong to any subset.

The array:

buckets = [0] * k

stores the current sum of each subset.

The recursive function places one number at a time:

def backtrack(index: int) -> bool:

If every number has been placed:

if index == len(nums):
    return True

then all buckets must equal target, because the total sum is exactly k * target and no bucket was allowed to exceed target.

For each bucket, we try placing the current value:

if buckets[bucket_index] + value > target:
    continue

If it fits, place it:

buckets[bucket_index] += value

Then recurse.

If it fails, undo the placement:

buckets[bucket_index] -= value

The symmetry pruning is:

if buckets[bucket_index] == 0:
    break

If placing the value into one empty bucket did not work, placing it into another empty bucket is equivalent.

Testing

def run_tests():
    s = Solution()

    assert s.canPartitionKSubsets(
        [4, 3, 2, 3, 5, 2, 1],
        4,
    ) is True

    assert s.canPartitionKSubsets(
        [1, 2, 3, 4],
        3,
    ) is False

    assert s.canPartitionKSubsets(
        [1, 1, 1, 1],
        2,
    ) is True

    assert s.canPartitionKSubsets(
        [2, 2, 2, 2, 3, 4, 5],
        4,
    ) is False

    assert s.canPartitionKSubsets(
        [10, 10, 10, 7, 7, 7, 7, 7, 7, 6, 6, 6],
        3,
    ) is True

    print("all tests passed")

run_tests()

Test meaning: