Problem description:

Given an array of integers nums and a positive integer k, find whether it’s possible to divide this array into k non-empty subsets whose sums are all equal.

Example 1:
Input: nums = [4, 3, 2, 3, 5, 2, 1], k = 4
Output: True
Explanation: It’s possible to divide it into 4 subsets (5), (1, 4), (2,3), (2,3) with equal sums.
Note:

1 <= k <= len(nums) <= 16.
0 < nums[i] < 10000.

Solution:

Assume sum is the sum of nums[] . The dfs process is to find a subset of nums[] which sum equals to sum/k. We use an array visited[] to record which element in nums[] is used. Each time when we get a cur_sum = sum/k, we will start from position 0 in nums[] to look up the elements that are not used yet and find another cur_sum = sum/k.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
class Solution:
    def canPartitionKSubsets(self, nums: List[int], k: int) -> bool:
        # backtracking
        # since we don't have to print the result, we could use visited array to note if an element is used in a subset
        
        total = sum(nums)
        if total == 0 or total % k != 0:
            return False
        visited = [0]* len(nums)
        
        
        def backtrack(visited, nums, k, cur_sum, cur_num, start, target):
            if k == 1: 
                # we call backtrack when it's able to make k subsets
                # when there're k-1 subsets existed, we don't have to calculate the last
                return True
            if cur_sum == target and cur_num > 0:
                # successfully find a subset
                return backtrack(visited, nums, k-1, 0, 0, 0, target)
            for i in range(start, len(nums)):
                if not visited[i]:
                    visited[i] = 1
                    if backtrack(visited, nums, k, cur_sum+nums[i], cur_num+1, i+1, target):
                        return True
                    visited[i] = 0
            return False
        return backtrack(visited, nums, k, 0, 0, 0, total/k)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
public:
    bool canPartitionKSubsets(vector<int>& nums, int k) {
        int sum= accumulate(nums.begin(), nums.end(), 0);
        if(sum % k != 0) return false; //can not form k subsets
        
        vector<bool> visited(nums.size(), false);
        return helper(nums, k, sum/k, 0, 0, visited);
    }
    
    bool helper(vector<int>& nums, int k, int target, int start, int curSum, vector<bool>& visited){
        if(k == 1) return true; //only need to create one subset, current subset can do it
        if(curSum == target) return helper(nums, k-1, target, 0, 0, visited); //achieve target, need to create k-1 more subset
        for(int i= start; i< nums.size(); i++){
            if(visited[i]) continue;
            visited[i]= true;
            if(helper(nums, k, target, i+1, curSum+nums[i], visited)) 
                return true;
            
            visited[i]= false;
        }
        return false;
    }
    
};