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

1
from functools import cache
2

3

4
class Solution:
5
    def numberWays(self, hats: List[List[int]]) -> int:
6
        # transpose hats
7
        n_persons = len(hats)
8
        temp = [[] for _ in range(40)]
9
        for i, person in enumerate(hats):
10
            for hat in person:
11
                temp[hat - 1].append(i)
12
        hats = temp
13

14
        # drop empty hats and sort
15
        hats = [h for h in hats if h]
16
        hats.sort(key=lambda h: len(h))
17
        n_hats = len(hats)
18

19
        # helpers
20
        full = 2**n_persons - 1
21
        bits = [1 << i for i in range(n_persons)]
22
        cut = 10**9 + 7
23

24
        @cache
25
        def get_count(i, has_hat):
26
            if has_hat == full:
27
                # all persons have a hat
28
                return 1
29
            elif i == n_hats:
30
                # no more hats left
31
                return 0
32
            else:
33
                result = 0
34
                for person in hats[i]:
35
                    # does the person
36
                    # still not have a hat?
37
                    if has_hat & bits[person] == 0:
38
                        # the person wears i-th hat
39
                        result += get_count(i + 1, has_hat | bits[person])
40
                # no one wears i-th hat
41
                result += get_count(i + 1, has_hat)
42

43
                # return modulo (10 ** 9 - 7)
44
                return result % cut
45

46
        # start form 0-th hat and
47
        # no person wearing a hat
48
        return get_count(0, 0)