Number of Ways to Wear Different Hats to Each Other LeetCode solution in C++

1
class Solution {
2
 private:
3
  int mod = 1e9 + 7;
4
  // stores the hat as the key and an array of the people who prefer it
5
  unordered_map<int, vector<int>> hatTable;
6
  // dp table defined by two variables:-
7
  // hat->the current hat number we are on
8
  // mask->is a bitmask which tells us which people have already been given a
9
  // hat
10
  vector<vector<int>> mem;
11
  // every hat may or may not be paired with a person, therefore for every hat
12
  // present in the hatTable hashmap we will either ignore it(not assign it to
13
  // anybody) or we will iterate over all the people who have a preference for
14
  // that hat and assign it to them. Then perform a recursive call for the next
15
  // hat with the updated bitmask and store the results.
16
  int solve(int hat, int mask, int n) {
17
    // if all the people have been assigned a hat, then we have discovered one
18
    // possible assignment
19
    if (mask == (1 << n) - 1) {
20
      return 1;
21
    } else if (hat == 41) {
22
      // otherwise if we have exhausted all the hats then return 0 as we were
23
      // not able to discover a valid assignment
24
      return 0;
25
    } else if (mem[hat][mask] != -1) {
26
      return mem[hat][mask];
27
    } else if (!hatTable.count(hat)) {
28
      return mem[hat][mask] = solve(hat + 1, mask, n) % mod;
29
    } else {
30
      // we can skip the current hat(not assign it to anybody)
31
      int ans = solve(hat + 1, mask, n);
32
      // we will assign the current hat to every person who prefers it and then
33
      // perform a recursive call and store the number of ways from all possible
34
      // assignments
35
      for (int i = 0; i < hatTable[hat].size(); i++) {
36
        int curr = hatTable[hat][i];
37
        if ((mask >> curr) & 1) {
38
          continue;
39
        } else {
40
          ans = (ans % mod + solve(hat + 1, mask | (1 << curr), n) % mod) % mod;
41
        }
42
      }
43
      return mem[hat][mask] = ans;
44
    }
45
  }
46

47
 public:
48
  int numberWays(vector<vector<int>> &hats) {
49
    int n = hats.size();
50
    for (int i = 0; i < hats.size(); i++) {
51
      for (int j = 0; j < hats[i].size(); j++) {
52
        hatTable[hats[i][j]].push_back(i);
53
      }
54
    }
55
    mem = vector<vector<int>>(41, vector<int>(1 << n, -1));
56
    return solve(1, 0, n);
57
  }
58
};