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
| class Solution:
def getMaxGridHappiness(self, rows, cols, introvertsCount, extrovertsCount):
introvertGain = 120
extrovertGain = 40
introvertLoss = -30
extrovertLoss = 20
neighborHappinessLoss = [[0, 0, 0], [0, 2 * introvertLoss, introvertLoss + extrovertLoss], [0, introvertLoss + extrovertLoss, 2 * extrovertLoss]]
@lru_cache(None)
def dp(currentIndex, currentRowState, remainingIntroverts, remainingExtroverts):
if currentIndex == -1:
return 0
currentRow = currentIndex // cols
currentCol = currentIndex % cols
answers = []
neighbors = [(1, remainingIntroverts - 1, remainingExtroverts, introvertGain),
(2, remainingIntroverts, remainingExtroverts - 1, extrovertGain),
(0, remainingIntroverts, remainingExtroverts, 0)]
for val, newIntroverts, newExtroverts, gain in neighbors:
currentHappiness = 0
if newIntroverts >= 0 and newExtroverts >= 0:
currentHappiness = dp(currentIndex - 1, (val,) + currentRowState[:-1], newIntroverts, newExtroverts) + gain
if currentCol < cols - 1: currentHappiness += neighborHappinessLoss[val][currentRowState[0]] # right neighbor
if currentRow < rows - 1: currentHappiness += neighborHappinessLoss[val][currentRowState[-1]] # down neighbor
answers.append(currentHappiness)
return max(answers)
if rows < cols:
rows, cols = cols, rows
return dp(rows * cols - 1, tuple([0] * cols), introvertsCount, extrovertsCount)
|