Maximize Grid Happiness LeetCode solution in C++

1
class Solution {
2
 public:
3
  int get_ternary_mask_bit(int mask, int i) {
4
    mask /= (int)pow(3, i);
5
    return mask % 3;
6
  }
7
  int get_next_mask(int mask, int val, int n) {
8
    return (mask * 3 + val) % (243);
9
  }
10
  int getMaxGridHappiness(int m, int n, int introvertsCount, int extrovertsCount) {
11
    // --> index introvert_left extrovert_left introvert_mask extrovert_mask
12
    // 0 -> empty
13
    // 1 -> introvert
14
    // 2 -> extrovert
15
    vector<vector<vector<vector<int>>>> dp(
16
        26, vector<vector<vector<int>>>(7, vector<vector<int>>(7, vector<int>(244, INT_MIN))));
17
    if (m > n) swap(m, n);
18
    if (m == 1 && n == 1) {
19
      if (introvertsCount)
20
        return 120;
21
      else if (extrovertsCount)
22
        return 40;
23
      else
24
        return 0;
25
    }
26
    map<pair<int, int>, int> score_map_intro = {{{0, 0}, 120},
27
                                                {{0, 1}, 120 - 30 - 30},
28
                                                {{0, 2}, 120 + 20 - 30},
29
                                                {{1, 1}, 120 - 30 - 30 - 30 - 30},
30
                                                {{1, 2}, 120 - 30 - 30 - 30 + 20},
31
                                                {{2, 2}, 120 - 30 - 30 + 20 + 20}};
32
    map<pair<int, int>, int> score_map_extro = {{{0, 0}, 40},
33
                                                {{0, 1}, 40 - 30 + 20},
34
                                                {{0, 2}, 40 + 20 + 20},
35
                                                {{1, 1}, 40 + 20 + 20 - 30 - 30},
36
                                                {{1, 2}, 40 + 20 + 20 + 20 - 30},
37
                                                {{2, 2}, 40 + 20 + 20 + 20 + 20}};
38
    int ans = 0;
39
    for (int a = 0; a < m * n; a++) {
40
      for (int b = introvertsCount; b >= 0; b--) {
41
        for (int c = extrovertsCount; c >= 0; c--) {
42
          for (int mask = 0; mask < pow(3, 5); mask++) {
43
            int i = a / n;
44
            int j = a % n;
45
            int up_bit = i > 0 ? (get_ternary_mask_bit(mask, n - 1)) : (0);
46
            int left_bit = j > 0 ? (get_ternary_mask_bit(mask, 0)) : (0);
47
            if (a == 0) dp[a][b][c][mask] = 0;
48
            if (left_bit > up_bit) swap(left_bit, up_bit);
49
            // case 1 making this cell empty
50
            dp[a + 1][b][c][get_next_mask(mask, 0, n)] =
51
                max(dp[a + 1][b][c][get_next_mask(mask, 0, n)], dp[a][b][c][mask]);
52
            // case 2 making this cell introvert
53
            if (b > 0)
54
              dp[a + 1][b - 1][c][get_next_mask(mask, 1, n)] =
55
                  max(dp[a + 1][b - 1][c][get_next_mask(mask, 1, n)],
56
                      dp[a][b][c][mask] + score_map_intro[{left_bit, up_bit}]);
57
            // case 3 making this cell extrovert
58
            if (c > 0)
59
              dp[a + 1][b][c - 1][get_next_mask(mask, 2, n)] =
60
                  max(dp[a + 1][b][c - 1][get_next_mask(mask, 2, n)],
61
                      dp[a][b][c][mask] + score_map_extro[{left_bit, up_bit}]);
62
          }
63
        }
64
      }
65
    }
66
    for (int b = introvertsCount; b >= 0; b--)
67
      for (int c = extrovertsCount; c >= 0; c--)
68
        for (int mask = 0; mask < pow(3, 5); mask++) ans = max(ans, dp[n * m][b][c][mask]);
69

70
    return ans;
71
  }
72
};