Maximize Grid Happiness LeetCode solution in Java

1
class Solution {
2
  public int getMaxGridHappiness(int m, int n, int introvertsCount, int extrovertsCount) {
3
    // introvert = -1, extrovert = 1
4
    return dfs(
5
        m,
6
        n,
7
        0,
8
        introvertsCount,
9
        extrovertsCount,
10
        0,
11
        0,
12
        new Integer[m + 1][introvertsCount + 1][extrovertsCount + 1][(1 << (n + 1))]
13
            [(1 << (n + 1))]);
14
  }
15

16
  public int dfs(
17
      int m,
18
      int n,
19
      int row,
20
      int in_count,
21
      int ex_count,
22
      int prev_in_pos,
23
      int prev_ex_pos,
24
      Integer[][][][][] dp) {
25
    if (dp[row][in_count][ex_count][prev_in_pos][prev_ex_pos] != null)
26
      return dp[row][in_count][ex_count][prev_in_pos][prev_ex_pos];
27
    if ((in_count == 0 && ex_count == 0) || (row == m)) return 0;
28
    List<int[]> possible_permutations =
29
        new ArrayList<
30
            int[]>(); // get all possible ways to fill a row**** with given number of introverts &
31
    // extroverts
32
    int[] aux = new int[n];
33
    getPermutations(in_count, ex_count, aux, 0, possible_permutations, n);
34

35
    int ans = 0;
36
    for (int[] possible : possible_permutations) {
37
      int curr_in_count = in_count, curr_ex_count = ex_count;
38
      int curr_in_pos = 0, curr_ex_pos = 0;
39
      for (int i = 0; i < n; i++) {
40
        if (possible[i] == 0) continue;
41
        if (possible[i] == -1) {
42
          curr_in_count--;
43
          curr_in_pos |= (1 << i);
44
        } else {
45
          curr_ex_count--;
46
          curr_ex_pos |= (1 << i);
47
        }
48
      }
49
      int curr_row_val = calculate(possible, prev_in_pos, prev_ex_pos, n);
50
      int rest_rows_val =
51
          dfs(m, n, row + 1, curr_in_count, curr_ex_count, curr_in_pos, curr_ex_pos, dp);
52
      ans = Math.max(ans, curr_row_val + rest_rows_val);
53
    }
54
    return dp[row][in_count][ex_count][prev_in_pos][prev_ex_pos] = ans;
55
  }
56

57
  // for each row, find happiness (keeping in account : left, right and upper neighors, for people
58
  // in this row) and
59
  //                              (keeping in account : lower neighbors for people in previous row)
60
  public int calculate(int[] currRow, int prev_in_pos, int prev_ex_pos, int columns) {
61
    int res = 0;
62
    // vertical neighbors
63
    for (int i = 0; i < columns; i++) {
64
      if (currRow[i] == 0) continue;
65
      if (currRow[i] == 1) {
66
        res += 40;
67
        if ((prev_in_pos & (1 << i)) != 0)
68
          res += (20 - 30); // -30 : because previous upper neighbor is introvert
69
        else if ((prev_ex_pos & (1 << i)) != 0) res += (20 + 20);
70
      } else {
71
        res += 120;
72
        if ((prev_in_pos & (1 << i)) != 0) res += (-30 - 30);
73
        else if ((prev_ex_pos & (1 << i)) != 0) res += (-30 + 20);
74
      }
75
    }
76
    // horizontal neighbors
77
    for (int i = 0; i < columns; i++) {
78
      if (currRow[i] == 0) continue;
79
      if (currRow[i] == -1) {
80
        if (i - 1 >= 0 && currRow[i - 1] != 0) res += (-30);
81
        if (i + 1 < columns && currRow[i + 1] != 0) res += (-30);
82
      } else {
83
        if (i - 1 >= 0 && currRow[i - 1] != 0) res += (20);
84
        if (i + 1 < columns && currRow[i + 1] != 0) res += (20);
85
      }
86
    }
87
    return res;
88
  }
89

90
  public void getPermutations(
91
      int in_count,
92
      int ex_count,
93
      int[] curr,
94
      int index,
95
      List<int[]> possible_permutations,
96
      int len) {
97
    if ((in_count == 0 && ex_count == 0) || index == len) {
98
      int[] arr = new int[len];
99
      for (int i = 0; i < len; i++) arr[i] = curr[i];
100
      possible_permutations.add(arr);
101
      return;
102
    }
103
    if (in_count > 0) {
104
      curr[index] = -1;
105
      getPermutations(in_count - 1, ex_count, curr, index + 1, possible_permutations, len);
106
      curr[index] = 0;
107
    }
108
    if (ex_count > 0) {
109
      curr[index] = 1;
110
      getPermutations(in_count, ex_count - 1, curr, index + 1, possible_permutations, len);
111
      curr[index] = 0;
112
    }
113
    getPermutations(in_count, ex_count, curr, index + 1, possible_permutations, len);
114
  }
115
}