The trick here is to name the state correctly, then let the implementation follow. For 698. Partition to K Equal Sum Subsets, the solution in this repository is mainly a DFS + memoization solution.
Guide
What?
We want to turn the problem statement into a smaller set of decisions the computer can repeat safely. Instead of trying to be clever immediately, read the code as a sequence of questions:
- What state are we keeping?
- How do we move from one state to the next?
- When do we know the answer is already determined?
For this file, the main tools are: DFS + memoization, dynamic programming, backtracking.
The notes already sitting in the source point us in the right direction:
- dfs
- TLE
- 64 / 149 test cases passed.
Guide
When?
This pattern shows up when the brute force version has too many repeated checks, too many possible branches, or too much bookkeeping to do by hand. The accepted code reduces that pressure by storing exactly the information that remains useful later.
The important function names to track are backtrack, canPartitionKSubsets.
Guide
Why?
The solution works because it narrows the problem until every update has a clear reason to exist.
- Sorting is used to make local choices comparable, so the later scan does not have to rediscover order.
- The final return is not magic; it is the invariant after the loops or recursion have finished doing their accounting.
Guide
How?
Walk through the solution in this order:
- Read the setup variables first.
- Follow the main loop or recursive helper next.
- Watch where invalid states get skipped.
- Check which value survives to the return statement.
The most important competitive-programming habit here is to trust the invariant. Once the invariant is right, the loops become much less scary.
Guide
Complexity
- Time: O(k^(N-k)*k!)
- Space: O(N)
Guide
C++ Solution
Your submission
The accepted solution
solution.cpp
01//dfs
02//TLE
03//64 / 149 test cases passed.
04class Solution {
05public:
06 int partitionSum;
07 vector<int> subsets;
08 bool ans;
09
10 void backtrack(vector<int>& nums, int k, int start){
11 if(start == nums.size()){
12 for(int val : subsets){
13 cout << val << " ";
14 }
15 cout << endl;
16 ans = all_of(subsets.begin(), subsets.end(),
17 [this](const int& val){return val == partitionSum;});
18 }else if(!ans){ //early stop
19 for(int i = 0; i < k; i++){
20 //early stop
21 if(subsets[i] + nums[start] > partitionSum)
22 continue;
23 subsets[i] += nums[start];
24 backtrack(nums, k, start+1);
25 subsets[i] -= nums[start];
26 }
27 }
28 };
29
30 bool canPartitionKSubsets(vector<int>& nums, int k) {
31 subsets = vector<int>(k, 0);
32 int sum = accumulate(nums.begin(), nums.end(), 0);
33 if(sum % k != 0) return false;
34 partitionSum = sum/k;
35 ans = false;
36 backtrack(nums, k, 0);
37 return ans;
38 }
39};
40
41//dfs, add some tricks to speed up
42//Runtime: 0 ms, faster than 100.00% of C++ online submissions for Partition to K Equal Sum Subsets.
43//Memory Usage: 9.4 MB, less than 40.17% of C++ online submissions for Partition to K Equal Sum Subsets.
44//time: O(k^(N-k)*k!)
45//space: O(N)
46class Solution {
47public:
48 int partitionSum;
49 vector<int> subsets;
50 bool ans;
51
52 void backtrack(vector<int>& nums, int k, int start){
53 if(start == nums.size()){
54 ans = all_of(subsets.begin(), subsets.end(),
55 [this](const int& val){return val == partitionSum;});
56 }else if(!ans){ //early stop
57 for(int i = 0; i < k; i++){
58 //early stop
59 if(subsets[i] + nums[start] > partitionSum)
60 continue;
61 subsets[i] += nums[start];
62 backtrack(nums, k, start+1);
63 subsets[i] -= nums[start];
64 /*
65 https://leetcode.com/problems/partition-to-k-equal-sum-subsets/solution/
66 speed up: TLE -> 584ms
67 if nums[start] does not fit an empty group,
68 (subsets[i] == 0 means an empty group),
69 it doesn't fit the following other empty groups too
70 when subsets[i] == 0, subsets[i+1] will always be 0
71 */
72 if(subsets[i] == 0) break;
73 }
74 }
75 };
76
77 bool canPartitionKSubsets(vector<int>& nums, int k) {
78 int sum = accumulate(nums.begin(), nums.end(), 0);
79 if(sum % k != 0) return false;
80 partitionSum = sum/k;
81 ans = false;
82 /*
83 speed up: 584ms -> 4ms
84 start to build subset from larger elements first,
85 so that the smaller subset can be considered earlier
86 */
87 sort(nums.rbegin(), nums.rend());
88 //speed up (0th element is the largest)
89 if(nums[0] > partitionSum) return false;
90 //speed up : remove the elements just equal to partitionSum
91 //the two speed up above: 4ms -> 0ms
92 int start = 0;
93 while(start < nums.size() &&
94 nums[start] == partitionSum){
95 start++;
96 k--;
97 }
98 subsets = vector<int>(k, 0);
99 backtrack(nums, k, start);
100 return ans;
101 }
102};
103
104//dp + bitmask
105//https://leetcode.com/problems/partition-to-k-equal-sum-subsets/discuss/480707/C%2B%2B-DP-bit-manipulation-in-20-lines
106//Runtime: 88 ms, faster than 26.04% of C++ online submissions for Partition to K Equal Sum Subsets.
107//Memory Usage: 17.4 MB, less than 12.91% of C++ online submissions for Partition to K Equal Sum Subsets.
108class Solution {
109public:
110 bool canPartitionKSubsets(vector<int>& nums, int k) {
111 int n = nums.size();
112 int sum = accumulate(nums.begin(), nums.end(), 0);
113 if(sum % k != 0) return false;
114 int partitionSum = sum/k;
115
116 //len(nums) <= 16, so there are 2^16 states
117 vector<int> dp(1<<n, -1);
118 dp[0] = 0;
119
120 //each mask represents a state
121 for(int mask = 0; mask < (1 << n); mask++){
122 //skip invalid state
123 if(dp[mask] == -1) continue;
124 //try to add nums[i] into the state
125 for(int i = 0; i < n; i++){
126 /*
127 dp[mask]: (sum of all subsets) % partitionSum
128 dp[mask] + nums[i] <= partitionSum: we must choose a number
129 that can be fit into a subset
130 */
131 if(!(mask & (1<<i)) && dp[mask] + nums[i] <= partitionSum){
132 /*
133 we set the dp value as xxx % partitionSum,
134 so that next time we can choose a number that can fit
135 in a subset by checking
136 dp[mask] + nums[i] <= partitionSum
137 */
138 dp[mask | (1<<i)] = (dp[mask] + nums[i]) % partitionSum;
139 }
140 }
141 }
142
143 //it's importatnt to add () around 1<<i !!
144 //'-' 's precedence is higher than '<<' !!
145 return dp[(1<<n) - 1] == 0;
146 }
147};Cost
Complexity
Time
O(k^(N-k)*k!)
Dominated by the main traversal, recursion, or data-structure operations in the code.
Space
O(N)
Auxiliary state plus the answer structure where the problem requires one.