diff --git a/leetcode/stuff.py b/leetcode/stuff.py index 90dc403..038c2d0 100644 --- a/leetcode/stuff.py +++ b/leetcode/stuff.py @@ -1,3 +1,6 @@ +import enum +import itertools +import pprint import typing @@ -48,3 +51,81 @@ def find_bisect(lower: float, upper: float, check: typing.Callable[[float], int] return find_bisect(mid, upper, check) return find_bisect(lower, mid, check) + + +def cartesian_path(p0: tuple[int, int], p1: tuple[int, int]) -> list[tuple[int, int]]: + path: list[tuple[int, int]] = [] + + if p0 < p1: + for i in range(p0[1], p1[1]): + path.append((i, p0[0])) + + for i in range(p0[0], p1[0]): + path.append((p1[1], i)) + + else: + for i in range(p0[1], p1[1] - 1, -1): + path.append((i, p0[0])) + + for i in range(p0[0] - 1, p1[0], -1): + path.append((p1[1], i)) + + return path + + +def gen_matrix(width: int, height: int) -> list[list[int]]: + return [list(range(width)) for _ in range(height)] + + +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 + + +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( + [ + (1, 0), # east + (0, 1), # south + (-1, 0), # west + (0, -1), # north + ] + ) + 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 > 3: + return False + + self.missteps += 1 + return self.step() + + self.missteps: int = 0 + self.path.append(next_loc) + return True diff --git a/leetcode/test_stuff.py b/leetcode/test_stuff.py index 1820bac..c10dda5 100644 --- a/leetcode/test_stuff.py +++ b/leetcode/test_stuff.py @@ -16,3 +16,46 @@ import stuff ) def test_find_sqrt_ish(n: int, expected: int): assert stuff.find_sqrt_ish(n) == expected + + +@pytest.mark.parametrize( + ("matrix", "expected"), + [ + ( + [ + ["a", "b", "c"], + ["d", "e", "f"], + ["g", "h", "i"], + ], + ["a", "b", "c", "f", "i", "h", "g", "d", "e"], + ), + ( + [ + [1, 2, 3, 4], + [5, 6, 7, 8], + [9, 10, 11, 12], + ], + [1, 2, 3, 4, 8, 12, 11, 10, 9, 5, 6, 7], + ), + ( + [ + [1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 2, 3, 4, 5, 6, 7, 8, 9], + ], + [] + + [1, 2, 3, 4, 5, 6, 7, 8, 9] # right + + [9, 9, 9] # down + + [9, 8, 7, 6, 5, 4, 3, 2, 1] # left + + [1, 1] # up + + [1, 2, 3, 4, 5, 6, 7, 8] # right + + [8] # down + + [8, 7, 6, 5, 4, 3, 2] # left + + [2, 3, 4, 5, 6, 7], # right + ), + ], +) +def test_matrix_spiral(matrix, expected): + assert stuff.matrix_spiral(matrix) == expected