Partition to K Equal Sum Subsets — Bitmask DP
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