Super-path that takes all grid cells to a sink cell

paths.py

neigh = {
    'R': (0, 1, 'L'), 'L': (0, -1, 'R'), 'D': (1, 0, 'U'), 'U': (-1, 0, 'D')
}


# Type of a cell
# -1: out of bounds. Otherwise the value of the cell
def cell_type(grid: list[list[int]], i: int, j: int) -> int:
    if 0 <= i < len(grid) and 0 <= j < len(grid[0]):
        return grid[i][j]
    return -1


# Shortest paths from all empty (0-valued) cells to sink (2-valued) cells
# Found by a combined reverse bfs from all the sink cells
def shortest_paths(grid: list[list[int]]) -> list[list[str]]:
    m, n = len(grid), len(grid[0])
    shortest = [[''] * n for _ in range(m)]
    bfsq = [(i, j) for i in range(m) for j in range(n) if grid[i][j] == 2]
    while bfsq:
        next = []
        for i, j in bfsq:
            for dx, dy, d in neigh.values():
                x, y = i + dx, j + dy
                if cell_type(grid, x, y) == 0 and not shortest[x][y]:
                    shortest[x][y] = d + shortest[i][j]
                    next.append((x, y))
        bfsq = next
    return shortest


# Find end cell of a path starting at (i, j) and following moves
# Move is ignored if it leads to a wall or out of bounds
# Move stops if it reaches a sink cell
def move(grid: list[list[int]], i: int, j: int, moves: str) -> tuple[int, int]:
    for m in moves:
        dx, dy, _ = neigh[m]
        ii, jj = i + dx, j + dy
        cell = cell_type(grid, ii, jj)
        if cell == 2:
            return ii, jj
        elif cell == 0:
            i, j = ii, jj
    return i, j


# Find a string that allows to reach a sink by starting from any empty cell
def super_path(grid: list[list[int]]) -> str:
    shortest = shortest_paths(grid)

    # All empty cells we need to move to sink cells
    todo = set([
        (i, j) for i in range(len(grid)) for j in range(len(grid[0]))
        if grid[i][j] == 0
    ])

    ret = ''
    while todo:
        next = set()

        # Find the shortest path among any remaining vertex to a sink cell
        best = ''
        for b, (i, j) in enumerate(todo):
            if b == 0 or len(shortest[i][j]) < len(best):
                best = shortest[i][j]

        # Move all vertices along the best path
        ret += best

        # Find where the candidate vertices end up after following the path
        # Retain all unique non-sink end points
        for i, j in todo:
            ii, jj = move(grid, i, j, best)
            if grid[ii][jj] != 2:
                next.add((ii, jj))

        # And continue the process with these
        todo = next
    return ret


print(super_path([[1, 0, 1], [0, 2, 0], [1, 0, 1]]))
print(super_path([[0, 0, 0], [0, 1, 0], [0, 2, 0]]))