Maximum Number of Groups Getting Fresh Donuts LeetCode solution in Python

1
class Solution:
2
    def maxHappyGroups(self, B, groups):
3
        ans = sum(g % B == 0 for g in groups)
4
        groups = [g for g in groups if g % B != 0]
5

6
        pos = [0] * B
7
        for g in groups:
8
            pos[g % B] += 1
9

10
        for i in range(1, B):
11
            t = min(pos[i], pos[B - i]) if 2 * i != B else pos[i] // 2
12
            ans += t
13
            pos[i] -= t
14
            pos[B - i] -= t
15

16
        if sum(pos) == 0:
17
            return ans
18

19
        @lru_cache(None)
20
        def dfs(position, last):
21
            if sum(position) == 0:
22
                return 0
23

24
            ans = float("-inf")
25
            for i in range(B):
26
                if position[i] > 0:
27
                    t = [j for j in position]
28
                    t[i] -= 1
29
                    U = (last - i) % B
30
                    ans = max(ans, dfs(tuple(t), U) + (U == 0))
31

32
            return ans
33

34
        return max(dfs(tuple(pos), i) for i in range(1, B)) + ans