3459. Length of Longest V-Shaped Diagonal Segment

from typing import List
from functools import cache

def lenOfVDiagonal(self, grid: List[List[int]]) -> int:
    # Get the dimensions of the grid
    n, m = len(grid), len(grid[0])
    # Define the four diagonal directions: down-right, down-left, up-left, up-right
    ds = ((1, 1), (1, -1), (-1, -1), (-1, 1))
    # Mapping for expected cell values:
    # When current expected value is 1, next becomes 2; when it is 2, next becomes 0 (termination)
    nx = [2, 2, 0]

    @cache
    def dp(i: int, j: int, x: int, d: int, turn_allowed: int) -> int:
        """
        Recursive function to compute the length of a diagonal sequence starting from (i, j)
        - i, j: current cell coordinates.
        - x: expected value at cell (i, j).
        - d: current direction index (0 to 3, corresponding to ds).
        - turn_allowed: flag indicating if a turn (clockwise) is allowed (1 means allowed, 0 means not allowed).
        
        Returns the length of the valid sequence from the current position.
        """
        # Check boundary conditions
        if not (0 <= i < n and 0 <= j < m):
            return 0
        # If the current cell does not have the expected value, stop recursion
        if grid[i][j] != x:
            return 0

        # Continue in the same direction without turning
        res = dp(i + ds[d][0], j + ds[d][1], nx[x], d, turn_allowed) + 1

        # If a turn is allowed, try turning clockwise (d2 is the next direction)
        if turn_allowed > 0:
            d2 = (d + 1) % 4
            res = max(res, dp(i + ds[d2][0], j + ds[d2][1], nx[x], d2, 0) + 1)
        return res

    res = 0
    # Start a new diagonal search from each cell that equals 1
    for i in range(n):
        for j in range(m):
            if grid[i][j] == 1:
                # Try all 4 directions starting from the current cell,
                # allowing one turn (turn_allowed = 1)
                res = max(res, max(dp(i, j, 1, d, 1) for d in range(4)))
    return res