Cherry Pickup II LeetCode solution in Python

1
class Solution:
2
    def cherryPickup(self, grid: List[List[int]]) -> int:
3
        rows, cols = len(grid), len(grid[0])
4

5
        dp = [[[0] * (cols + 2) for _ in range(cols + 2)] for _ in range(rows + 1)]
6

7
        def get_next_max(row, col_r1, col_r2):
8
            res = 0
9
            for next_col_r1 in (col_r1 - 1, col_r1, col_r1 + 1):
10
                for next_col_r2 in (col_r2 - 1, col_r2, col_r2 + 1):
11
                    res = max(res, dp[row + 1][next_col_r1 + 1][next_col_r2 + 1])
12

13
            return res
14

15
        for row in reversed(range(rows)):
16
            for col_r1 in range(min(cols, row + 2)):
17
                for col_r2 in range(max(0, cols - row - 1), cols):
18

19
                    reward = grid[row][col_r1] + grid[row][col_r2]
20
                    if col_r1 == col_r2:
21
                        reward /= 2
22

23
                    dp[row][col_r1 + 1][col_r2 + 1] = reward + get_next_max(
24
                        row, col_r1, col_r2
25
                    )
26

27
        return dp[0][1][cols]