LeetCodeInHardWay/src/Problem_0698_canPartitionKSubsets.cc at master · yinghaoyu/LeetCodeInHardWay · GitHub

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
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
#include <algorithm>
#include <vector>

#include "UnitTest.h"

using namespace std;

// 状压dp
// @sa https://www.bilibili.com/video/BV15a4y1o7NA/
class Solution
{
 public:
  // 记忆化搜索
  // 由于数组长1~16，state每1位bit表示数组元素有没有使用，二进制枚举
  // pre 表示前面元素的累加和
  // dp
  bool f(vector<int>& nums, int state, int pre, int avg, vector<bool>& dp)
  {
    if (state == 0)
    {
      // 所有元素都被使用
      return true;
    }
    if (!dp[state])
    {
      return dp[state];
    }
    dp[state] = false;  // 初始时只要state不是0，都是false
    for (int i = 0; i < nums.size(); i++)
    {
      if (nums[i] + pre > avg)
      {
        // 前面元素累加和 + 当前元素 > avg
        // nums已经有序，下一个元素一定更大，说明不用尝试了 !
        break;
      }
      if ((state >> i) & 1)  // 当前i位置的元素未使用
      {
        // 这里 % 的作用是，当 pre + num[i] == avg 的时候清0
        if (f(nums, state ^ (1 << i), ((pre + nums[i]) % avg), avg, dp))
        {
          dp[state] = true;
          break;
        }
      }
    }
    return dp[state];
  }

  bool canPartitionKSubsets(vector<int>& nums, int k)
  {
    int sum = 0;
    int N = nums.size();
    for (int i = 0; i < N; i++)
    {
      sum += nums[i];
    }
    int avg = sum / k;
    if (avg * k != sum)
    {
      return false;
    }
    std::sort(nums.begin(), nums.end());
    if (nums.back() > avg)
    {
      return false;
    }
    // dp[i]表示示在可用的数字状态为 i 的情况下是否可行
    // 这里的初值用false，结果没错，但会超时
    vector<bool> dp(1 << N, true);
    return f(nums, (1 << N) - 1, 0, avg, dp);
  }

  bool dp(vector<int>& nums, int k)
  {
    int sum = 0;
    int N = nums.size();
    for (int i = 0; i < N; i++)
    {
      sum += nums[i];
    }
    int avg = sum / k;
    if (avg * k != sum)
    {
      return false;
    }
    std::sort(nums.begin(), nums.end());
    if (nums.back() > avg)
    {
      return false;
    }
    int M = 1 << N;
    vector<int> dp(M, false);
    vector<int> pre(M, 0);
    dp[0] = true;                // 所有元素都被使用
    for (int i = 0; i < M; i++)  // state反向推
    {
      if (!dp[i])  // 剪枝，只有当前状态为true才会去推next状态
      {
        continue;
      }
      for (int j = 0; j < N; j++)  // 尝试
      {
        if (pre[i] + nums[j] > avg)  // 当前状态下，后面都不用尝试了
        {
          break;
        }
        if (((i >> j) & 1) == 0)  // i状态下，数组的第j元素已使用
        {
          int next = i | (1 << j);
          // 去推测next状态，这里跟dp[i]是有依赖关系的
          // 隐含条件： dp[i] == true
          if (!dp[next])
          {
            pre[next] = (pre[i] + nums[j]) % avg;
            dp[next] = true;
          }
        }
      }
    }
    return dp[M - 1];
  }
};

void testCanPartitionKSubsets()
{
  Solution s;
  vector<int> arr1 = {4, 3, 2, 3, 5, 2, 1};
  vector<int> arr2 = {1, 2, 3, 4};
  vector<int> arr3 = {2, 2, 2, 2, 3, 4, 5};
  vector<int> arr4 = {18, 20, 39, 73, 96, 99, 101, 111, 114, 190, 207, 295, 471, 649, 700, 1037};
  EXPECT_TRUE(s.canPartitionKSubsets(arr1, 4));
  EXPECT_FALSE(s.canPartitionKSubsets(arr2, 3));
  EXPECT_FALSE(s.canPartitionKSubsets(arr3, 4));
  EXPECT_TRUE(s.canPartitionKSubsets(arr4, 4));

  EXPECT_TRUE(s.dp(arr1, 4));
  EXPECT_FALSE(s.dp(arr2, 3));
  EXPECT_FALSE(s.dp(arr3, 4));
  EXPECT_TRUE(s.dp(arr4, 4));
  EXPECT_SUMMARY;
}

int main()
{
  testCanPartitionKSubsets();
  return 0;
}