Maximize Grid Happiness LeetCode solution in Python

1
class Solution:
2
    def getMaxGridHappiness(
3
        self, m: int, n: int, introvertsCount: int, extrovertsCount: int
4
    ) -> int:
5

6
        @cache
7
        def fn(prev, i, j, intro, extro):
8
            """Return max grid happiness at (i, j)."""
9
            if i == m:
10
                return 0  # no more position
11
            if j == n:
12
                return fn(prev, i + 1, 0, intro, extro)
13
            if intro == extro == 0:
14
                return 0
15

16
            prev0 = prev[:j] + (0,) + prev[j + 1 :]
17
            ans = fn(prev0, i, j + 1, intro, extro)
18
            if intro:
19
                val = 120
20
                if i and prev[j]:  # neighbor from above
21
                    val -= 30
22
                    if prev[j] == 1:
23
                        val -= 30
24
                    else:
25
                        val += 20
26
                if j and prev[j - 1]:  # neighbor from left
27
                    val -= 30
28
                    if prev[j - 1] == 1:
29
                        val -= 30
30
                    else:
31
                        val += 20
32
                prev0 = prev[:j] + (1,) + prev[j + 1 :]
33
                ans = max(ans, val + fn(prev0, i, j + 1, intro - 1, extro))
34
            if extro:
35
                val = 40
36
                if i and prev[j]:
37
                    val += 20
38
                    if prev[j] == 1:
39
                        val -= 30
40
                    else:
41
                        val += 20
42
                if j and prev[j - 1]:
43
                    val += 20
44
                    if prev[j - 1] == 1:
45
                        val -= 30
46
                    else:
47
                        val += 20
48
                prev0 = prev[:j] + (2,) + prev[j + 1 :]
49
                ans = max(ans, val + fn(prev0, i, j + 1, intro, extro - 1))
50
            return ans
51

52
        return fn((0,) * n, 0, 0, introvertsCount, extrovertsCount)