leetcode-go/medium/698.partition-to-k-equal-sum-subsets.go

type none struct{}

func canPartitionKSubsets(nums []int, k int) bool {
	sum := 0
	for _, i := range nums {
		sum += i
	}
	if sum%k != 0 {
		return false
	}
	sum /= k
	m := make(map[[3]int]none) // Use memorized search to avoid TLE
	return dfs(nums, k, sum, 0, 0, m)
}

func dfs(nums []int, k, sum, cur, set int, m map[[3]int]none) bool {
	if k == 0 {
		return set == 1<<uint(len(nums))-1
	}
	if _, ok := m[[3]int{k, cur, set}]; ok {
		return false
	}
	for i := range nums {
		if set>>uint(i)&1 == 1 {
			continue
		}
		if t := cur + nums[i]; sum < t {
			continue
		} else if t < sum {
			nset := set | (1 << uint(i))
			if dfs(nums, k, sum, t, nset, m) {
				return true
			}
			m[[3]int{k, t, nset}] = none{}
		} else {
			nset := set | (1 << uint(i))
			if dfs(nums, k-1, sum, 0, nset, m) {
				return true
			}
			m[[3]int{k, t, nset}] = none{}
		}
	}
	return false
}