LeetCode 1434.

Решал контест на LeetCode и встретил интересную задачку.



Моё решение:

 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
class Solution {
    private static final int MODULO = (int) 1e9 + 7;
    
    public int numberWays(List<List<Integer>> hats) {
        int nPeople = hats.size();
        boolean[][] preferred = new boolean[nPeople][41];
        for (int i = 0; i < nPeople; i++) {
            for (int hat : hats.get(i)) {
                preferred[i][hat] = true;
            }
        }
        return numberWays(40, preferred);
    }

    private int numberWays(int nHats, boolean[][] preferred) {
        int[][] dp = new int[nHats + 1][1 << preferred.length];
        for (int[] hatDp : dp) {
            Arrays.fill(hatDp, -1);
        }
        return numberWays(nHats, (1 << preferred.length) - 1, preferred, dp);
    }

    private int numberWays(int nHats, int gotHatMask, boolean[][] preferred, int[][] dp) {
        if (nHats == 0) {
            return gotHatMask == 0 ? 1 : 0;
        }
        if (dp[nHats][gotHatMask] == -1) {
            int sum = numberWays(nHats - 1, gotHatMask, preferred, dp);
            for (int i = 0; i < preferred.length; i++) {
                if (preferred[i][nHats] && getBit(gotHatMask, i) == 1) {
                    sum += numberWays(nHats - 1, clearBit(gotHatMask, i), preferred, dp);
                    sum %= MODULO;
                }
            }
            dp[nHats][gotHatMask] = sum;
        }
        return dp[nHats][gotHatMask];
    }

    private int getBit(int value, int nBit) {
        return (value >> nBit) & 1;
    }

    private int clearBit(int value, int nBit) {
        return value & ~(1 << nBit);
    }

}

Комментарии