box-o-sand/leetcode/spiral_matrix.py

import itertools
import typing


def matrix_spiral(matrix: list[list[typing.Any]]) -> list[typing.Any]:
    return [matrix[y][x] for x, y in matrix_spiral_path(matrix)]


def matrix_spiral_path(matrix: list[list[int]]) -> list[tuple[int, int]]:
    snek = SpinSnek(matrix)

    while snek.step():
        ...

    return snek.path


COMPASS = (
    (1, 0),  # east
    (0, 1),  # south
    (-1, 0),  # west
    (0, -1),  # north
)


class SpinSnek:
    def __init__(self, board: list[list[int]], loc: tuple[int, int] = (0, 0)):
        self.max_loc: tuple[int, int] = (len(board[0]) - 1, len(board) - 1)
        self.spinner: itertools.cycle[tuple[int, int]] = itertools.cycle(COMPASS)
        self.direction = next(self.spinner)
        self.path: list[tuple[int, int]] = [loc]
        self.missteps: int = 0

    def step(self) -> bool:
        loc = self.path[-1]
        next_loc: tuple[int, int] = (
            loc[0] + self.direction[0],
            loc[1] + self.direction[1],
        )

        if (
            next_loc[0] > self.max_loc[0]
            or next_loc[1] > self.max_loc[1]
            or next_loc[0] < 0
            or next_loc[1] < 0
            or next_loc in self.path
        ):
            self.direction = next(self.spinner)
            if self.missteps > len(COMPASS) - 1:
                return False

            self.missteps += 1
            return self.step()

        self.missteps: int = 0
        self.path.append(next_loc)
        return True