LeetCode 698. Partition to K Equal Sum Subsets

Post by ailswan Jul. 13, 2024

698. Partition to K Equal Sum Subsets

Problem Statement

link: LeetCode.cn LeetCode Given an integer array nums and an integer k, return true if it is possible to divide this array into k non-empty subsets whose sums are all equal.

Example:

Input: nums = [4,3,2,3,5,2,1], k = 4 Output: true

Input: nums = [1,2,3,4], k = 3 Output: false

Solution Approach

Algorithm

Implement

    class Solution:
    def canPartitionKSubsets(self, arr, k):
        total_sum = sum(arr)
        if total_sum % k != 0:tt
            return False
        
        target = total_sum // k
        return self.helper(arr, 0, k, target, 0, set())

    def helper(self, arr, curr, k, target, curr_selected_mask, selected_mask):
        if k == 1:
            return True
        if curr > target:
            return False

        if curr_selected_mask in selected_mask:
            return False

        if curr == target:
            res = self.helper(arr, 0, k - 1, target, curr_selected_mask, selected_mask)
            selected_mask.add(curr_selected_mask)
            return res

        for i in range(len(arr)):
            if (curr_selected_mask & (1 << i)) == 0:
                curr_selected_mask |= 1 << i

                if self.helper(arr, curr + arr[i], k, target, curr_selected_mask, selected_mask):
                    return True
                curr_selected_mask ^= 1 << i

        selected_mask.add(curr_selected_mask)
        return False