1
+ class Solution (object ):
2
+ def canPartitionKSubsets (self , nums , k ):
3
+ nums .sort (reverse = True )
4
+ buck , kSum = [0 ] * k , sum (nums ) // k
5
+
6
+ def dfs (idx ):
7
+ if idx == len (nums ):
8
+ return len (set (buck )) == 1 # every bucket sums to target
9
+ for i in range (k ):
10
+ buck [i ] += nums [idx ]
11
+ if buck [i ] <= kSum and dfs (idx + 1 ):
12
+ return True
13
+ buck [i ] -= nums [idx ]
14
+ if buck [i ] == 0 :
15
+ break
16
+ return False
17
+ return dfs (0 )
18
+
19
+
20
+ # Annotated and slower-but-clearer version of @jingkuan's solution based on @chemikadze's
21
+ # solution using @chengyuge925's solution.
22
+ class Solution :
23
+ def canPartitionKSubsets (self , nums , k ):
24
+ buckets = [0 ]* k
25
+ target = sum (nums ) // k
26
+
27
+ # We want to try placing larger numbers first
28
+ nums .sort (reverse = True )
29
+
30
+
31
+ # DFS determines which bucket to put the 'current element' (nums[idx] ) into
32
+ def dfs (idx ):
33
+ # If we've placed all of the items, we're done;
34
+ # check if we correctly made k equal subsets of
35
+ # size sum(nums) // k
36
+ if idx == len (nums ):
37
+ return set (buckets ) == set ([target ])
38
+
39
+ # For each bucket
40
+ for i in range (k ):
41
+ # Try adding the current element to it
42
+ buckets [i ] += nums [idx ]
43
+
44
+ # If it's a valid placement and we correctly placed the next element, we're
45
+ # done placing the current element.
46
+ if buckets [i ] <= target and dfs (idx + 1 ):
47
+ return True
48
+
49
+ # Otherwise, remove the current element from the ith bucket and
50
+ # try the next one.
51
+ buckets [i ] -= nums [idx ]
52
+
53
+ # This is an optimization that is not strictly necessary.
54
+ # If bucket[i] == 0, it means:
55
+ # - We put nums[idx] into an empty bucket
56
+ # - We tried placing every other element after and failed.
57
+ # - We took nums[idx] out of the bucket, making it empty again.
58
+ # So trying to put nums[idx] into a _different_ empty bucket will not produce
59
+ # a correct solution; we will just waste time (we place elements left to right,
60
+ # so if this bucket is now empty, every one after it is too).
61
+ #
62
+ # Otherwise (bucket[i] > 0), we just go to the next bucket and
63
+ # try placing nums[idx] there. If none of them work out, we wind up
64
+ # breaking out of the loop when range(k) ends and returning False.
65
+ if buckets [i ] == 0 :
66
+ break
67
+
68
+ # We couldn't place the current element anywhere that
69
+ # leads to a valid solution, so we will need to backtrack
70
+ # and try something else.
71
+ return False
72
+
73
+ # Start by trying to place nums[0]
74
+ return dfs (0 )
0 commit comments