Cherry Pickup Python Solution with Tests

Tested Python solution for LeetCode 741 with 14 pytest cases. Generate a practice environment with lcpy.

LeetCode 741, Hard. Topics: Array, Dynamic Programming, Matrix. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 14 parametrized pytest cases, and a playground notebook:

lcpy gen -n 741   # by problem number
lcpy gen -s cherry_pickup   # by problem name

​

Problem

You are given an n x n grid representing a field of cherries, each cell is one of three possible integers.

• 0 means the cell is empty, so you can pass through,
• 1 means the cell contains a cherry that you can pick up and pass through, or
• -1 means the cell contains a thorn that blocks your way.

Return the maximum number of cherries you can collect by following the rules below:

• Starting at the position (0, 0) and reaching (n - 1, n - 1) by moving right or down through valid path cells (cells with value 0 or 1).
• After reaching (n - 1, n - 1), returning to (0, 0) by moving left or up through valid path cells.
• When passing through a path cell containing a cherry, you pick it up, and the cell becomes an empty cell 0.
• If there is no valid path between (0, 0) and (n - 1, n - 1), then no cherries can be collected.

​

Examples

Input: grid = [[0,1,-1],[1,0,-1],[1,1,1]]
Output: 5
Explanation: The player started at (0, 0) and went down, down, right right to reach (2, 2).
4 cherries were picked up during this single trip, and the matrix becomes [[0,1,-1],[0,0,-1],[0,0,0]].
Then, the player went left, up, up, left to return home, picking up one more cherry.
The total number of cherries picked up is 5, and this is the maximum possible.

Input: grid = [[1,1,-1],[1,-1,1],[-1,1,1]]
Output: 0

​

Constraints

• n == grid.length
• n == grid[i].length
• 1 <= n <= 50
• grid[i][j] is -1, 0, or 1.
• grid[0][0] != -1
• grid[n - 1][n - 1] != -1

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

class Solution:
    # Time: O(n^3) steps x O(n^2) row pairs
    # Space: O(n^2)
    def cherry_pickup(self, grid: list[list[int]]) -> int:
        n = len(grid)
        # -1 marks unreachable states (cherry counts are always >= 0)
        unreachable = -1
        # dp[r1][r2]: max cherries with both walkers on diagonal r + c = t
        dp = [[unreachable] * n for _ in range(n)]
        dp[0][0] = grid[0][0]

        for t in range(1, 2 * n - 1):
            ndp = [[unreachable] * n for _ in range(n)]
            for r1 in range(max(0, t - n + 1), min(n, t + 1)):
                for r2 in range(r1, min(n, t + 1)):
                    c1, c2 = t - r1, t - r2
                    if grid[r1][c1] == -1 or grid[r2][c2] == -1:
                        continue
                    best = unreachable
                    for pr1 in (r1 - 1, r1):
                        for pr2 in (r2 - 1, r2):
                            if 0 <= pr1 < n and 0 <= pr2 < n and dp[pr1][pr2] > best:
                                best = dp[pr1][pr2]
                    if best < 0:
                        continue
                    cherries = grid[r1][c1]
                    if r1 != r2:
                        cherries += grid[r2][c2]
                    ndp[r1][r2] = best + cherries
            dp = ndp

        return max(dp[n - 1][n - 1], 0)

​

Complexity

Time	Space
O(n^3) steps x O(n^2) row pairs	O(n^2)

​

Tags

NeetCode All.