Bisecting and matrix goop

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

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