698. Partition to K Equal Sum Subsets

原本以為 DFS 寫這麼多年了,應該沒啥暴搜會寫爛,結果…

Sample code 裡面有這段代碼

1
2
if (sum[i] == 0)
    break;

想了半天,覺得太猛啦!他這個剪枝,可以把所有的空缺處都放在後段,也保證會連續出現在後面。

另外就是 sorting,由大到小。

時間複雜度會是 $O(1 * 2 * 3 * … * k * k^{n - k})$ = $O({k!} * k^{n - k})$,因為有上面的剪支,保證第一個數只能有一個可能的位子可以放,第二個數只有兩個…

AC Code

 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
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class Solution
{
private:
    vector<int> sum;
    bool dfs(int depth, int arr[], const vector<int> &nums, int k, int target)
    {
        if (depth == (int)nums.size()) {
            /*
            for (int i = 0; i < nums.size(); i++) {
                printf("%d %d\n", i, arr[i]);
            }
            */
            return true;
        }

        for (int i = 0; i < k; i++) {
            arr[depth] = i;
            sum[i] += nums[depth];
            
            if(sum[i] > target) {
                // skip
            } else {
                if (dfs(depth + 1, arr, nums, k, target))
                    return true;
            }
            
            sum[i] -= nums[depth];
            arr[depth] = -1;
            
            // we guarantee the sum[i] = 0 will be consecutive, and towards the end only!
            if(sum[i] == 0) // WTF?
                break;
        }

        return false;
    }

public:
    bool canPartitionKSubsets(vector<int> &nums, int k)
    {
        int tot = 0;
        for (auto i : nums)
            tot += i;
        if (tot % k != 0)
            return false;

        int target = tot / k;
        int arr[nums.size()];
        sum = vector<int>(k, 0);
        
        sort(nums.begin(), nums.end());
        reverse(nums.begin(), nums.end());
        for(auto i : nums) 
            if(i > target)
                return false;

        return dfs(0, arr, nums, k, target);
    }
};