Day 21: Step Counter (non-general solution)

Author: ephemient

README.md

@@ -25,5 +25,5 @@ Development occurs in language-specific directories:
 |[Day18.hs](hs/src/Day18.hs)|[Day18.kt](kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day18.kt)|[day18.py](py/aoc2023/day18.py)|[day18.rs](rs/src/day18.rs)|
 |[Day19.hs](hs/src/Day19.hs)|[Day19.kt](kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day19.kt)|[day19.py](py/aoc2023/day19.py)|[day19.rs](rs/src/day19.rs)|
 |[Day20.hs](hs/src/Day20.hs)|[Day20.kt](kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day20.kt)|[day20.py](py/aoc2023/day20.py)|[day20.rs](rs/src/day20.rs)|
-|[Day21.hs](hs/src/Day21.hs)|[Day21.kt](kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day21.kt)|||
+|[Day21.hs](hs/src/Day21.hs)|[Day21.kt](kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day21.kt)|[day21.py](py/aoc2023/day21.py)||
 |[Day22.hs](hs/src/Day22.hs)|[Day22.kt](kt/aoc2023-lib/src/commonMain/kotlin/com/github/ephemient/aoc2023/Day22.kt)|||

py/aoc2023/day21.py

@@ -0,0 +1,91 @@
+"""
+Day 21: Step Counter
+"""
+
+SAMPLE_INPUT = """
+...........
+.....###.#.
+.###.##..#.
+..#.#...#..
+....#.#....
+.##..S####.
+.##..#...#.
+.......##..
+.##.#.####.
+.##..##.##.
+...........
+"""
+
+
+def _count(grid, start, n):
+    frontier, visited, acc = {start}, set(), 0
+    for d in range(n):
+        if not (d ^ n) & 1:
+            acc += len(frontier)
+        visited |= frontier
+        frontier = {
+            (y1, x1)
+            for y0, x0 in frontier
+            for y1, x1 in [(y0 - 1, x0), (y0, x0 - 1), (y0, x0 + 1), (y0 + 1, x0)]
+            if 0 <= y1 < len(grid)
+            and 0 <= x1 < len(grid[y1])
+            and grid[y1][x1] != "#"
+            and (y1, x1) not in visited
+        }
+    return acc + len(frontier)
+
+
+def part1(data, n=64):
+    """
+    >>> part1(SAMPLE_INPUT, n=1)
+    2
+    >>> part1(SAMPLE_INPUT, n=2)
+    4
+    >>> part1(SAMPLE_INPUT, n=3)
+    6
+    >>> part1(SAMPLE_INPUT, n=6)
+    16
+    """
+    grid = [line for line in data.splitlines() if line]
+    (start,) = (
+        (y, x) for y, line in enumerate(grid) for x, c in enumerate(line) if c == "S"
+    )
+    return _count(grid, start, n)
+
+
+def part2(data, n=26501365):
+    """
+    >>> part2(SAMPLE_INPUT, n=6) # doctest: +SKIP
+    16
+    >>> part2(SAMPLE_INPUT, n=10) # doctest: +SKIP
+    50
+    >>> part2(SAMPLE_INPUT, n=50) # doctest: +SKIP
+    1594
+    >>> part2(SAMPLE_INPUT, n=100) # doctest: +SKIP
+    6536
+    >>> part2(SAMPLE_INPUT, n=500) # doctest: +SKIP
+    167004
+    >>> part2(SAMPLE_INPUT, n=1000) # doctest: +SKIP
+    668697
+    >>> part2(SAMPLE_INPUT, n=5000) # doctest: +SKIP
+    16733044
+    """
+    grid = [line for line in data.splitlines() if line]
+    m = len(grid)
+    q, r = n // m, n % m
+    ((y0, x0),) = (
+        (y, x) for y, line in enumerate(grid) for x, c in enumerate(line) if c == "S"
+    )
+    a, b, c, d = (
+        _count(
+            [line * (2 * i + 1) for line in grid] * (2 * i + 1),
+            (y0 + i * m, x0 + i * m),
+            r + i * m,
+        )
+        for i in range(4)
+    )
+    assert d == a - 3 * b + 3 * c
+    return a + (b - a) * q + (c - 2 * b + a) * (q * (q - 1) // 2)
+
+
+parts = (part1, part2)

py/pyproject.toml

@@ -42,6 +42,7 @@ day17 = "aoc2023.day17:parts"
 day18 = "aoc2023.day18:parts"
 day19 = "aoc2023.day19:parts"
 day20 = "aoc2023.day20:parts"
+day21 = "aoc2023.day21:parts"
 
 [tool.black]
 target_version = ["py312"]