Partition to K Equal Sum Subsets — Bitmask DP

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem

Determine if the array can be divided into k non-empty subsets each with equal sum.


Approach — Bitmask DP

dp[mask] = total sum placed into currently-filling bucket given elements in mask are used. Each element fills buckets in round-robin.

Time: O(2^n * n) | Space: O(2^n)


Solutions

Python — Bitmask DP

class Solution:
    def canPartitionKSubsets(self, nums: list[int], k: int) -> bool:
        total=sum(nums)
        if total%k: return False
        target=total//k
        n=len(nums); nums.sort(reverse=True)
        if nums[0]>target: return False
        dp=[-1]*(1<<n); dp[0]=0
        for mask in range(1<<n):
            if dp[mask]==-1: continue
            for i in range(n):
                if mask&(1<<i): continue
                next_mask=mask|(1<<i)
                if dp[mask]+nums[i]<=target:
                    dp[next_mask]=(dp[mask]+nums[i])%target
                    break
        return dp[(1<<n)-1]==0

Complexity

  • Time: O(2^n * n) | Space: O(2^n)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro