Maximum AND Sum of Array LeetCode solution in C++

1
class Solution {
2
 private:
3
  int dp[20][60000];
4
  vector<int> a;
5
  int n;
6
  int k;
7
  // n numbers, k slots
8
  int ans;
9
  int h[12]; // h[i] == 3^i
10
  int search(int start, int mask) {
11
    if (start == n) return 0;
12
    if (dp[start][mask] != -1) return dp[start][mask];
13
    // try to push a[start] into slot 1 ~ k
14
    for (int i = 1; i <= k; i++) {
15
      int cnt = (mask % h[i + 1]) / h[i];
16
      if (cnt <= 1) {
17
        dp[start][mask] = max(dp[start][mask], (a[start] & i) + search(start + 1, mask + h[i]));
18
      }
19
    }
20
    return dp[start][mask];
21
  }
22

23
 public:
24
  int maximumANDSum(vector<int> &nums, int numSlots) {
25
    // states: at most 3^k
26
    // search: k
27
    a.assign(nums.begin(), nums.end());
28
    memset(dp, -1, sizeof(dp));
29
    // dp[i][mask]: the answer for nums[i:] with mask   (based-3)
30
    // a 3^1 + b 3^2 + c 3^3 + ...
31
    n = nums.size();
32
    k = numSlots;
33
    ans = 0;
34
    h[0] = 1;
35
    for (int i = 1; i < 12; i++) h[i] = h[i - 1] * 3;
36
    return search(0, 0);
37
  }
38
};