day 21

aoc.sh

@@ -87,3 +87,6 @@ stack exec day19 < input/day19.input
 
 echo -e "${IYellow}--- Day 20: Pulse Propagation ---${Color_Off}"
 stack exec day20 < input/day20.input
+
+echo -e "${IRed}--- Day 21: Step Counter ---${Color_Off}"
+stack exec day21 < input/day21.input

aoc2023.cabal

@@ -459,6 +459,7 @@ test-suite spec
       TestDay18
       TestDay19
       TestDay20
+      TestDay21
       Paths_aoc2023
   autogen-modules:
       Paths_aoc2023

src/Day21.hs

@@ -1,10 +1,50 @@
 module Day21 (main, part1, part2) where
 
+import qualified Data.Set as S
+
+type Position = (Int, Int)
+type Grid = (S.Set Position, Int)
+
+coords :: String -> Char -> [Position]
+coords input char =
+    concatMap
+        ( \(r, line) ->
+            map (\(c, _) -> (r, c)) $
+                filter (\(_, c) -> c == char) $
+                    zip [0 ..] line
+        )
+        $ zip [0 ..]
+        $ lines input
+
+parse :: String -> (Position, Grid)
+parse input = (starting, (S.fromList $ starting : positions, n))
+  where
+    positions = coords input '.'
+    starting = head $ coords input 'S'
+    n = length $ lines input
+
+valid :: Grid -> Position -> Bool
+valid (ps, n) (r, c) = (r `mod` n, c `mod` n) `S.member` ps
+
+neighbors :: Position -> S.Set Position
+neighbors (r, c) = S.fromList [(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)]
+
+visit :: Grid -> S.Set Position -> S.Set Position
+visit grid = S.unions . S.map (S.filter (valid grid) . neighbors)
+
 part1 :: String -> String
-part1 = id
+part1 = show . length . go . parse
+  where
+    go (starting, grid) = iterate (visit grid) (S.singleton starting) !! 64
 
 part2 :: String -> String
-part2 = const ""
+part2 input = show $ case map (xs !!) $ take 3 [n `div` 2, n + n `div` 2 ..] of
+    [a, b, c] -> a + n' * (b - a) + n' * (n' - 1) `div` 2 * ((c - b) - (b - a))
+    _ -> error "unreachable"
+  where
+    (starting, grid@(_, n)) = parse input
+    xs = map length $ iterate (visit grid) (S.singleton starting)
+    n' = 26501365 `div` n
 
 solve :: String -> String
 solve input = "Part 1: " ++ part1 input ++ "\nPart 2: " ++ part2 input ++ "\n"

test/Main.hs

@@ -20,6 +20,7 @@ import TestDay17
 import TestDay18
 import TestDay19
 import TestDay20
+import TestDay21
 
 main :: IO ()
 main = do
@@ -43,3 +44,4 @@ main = do
     TestDay18.test
     TestDay19.test
     TestDay20.test
+    TestDay21.test

test/TestDay21.hs

@@ -0,0 +1,22 @@
+module TestDay21 (test) where
+
+import Day21
+import Test.Hspec
+
+input :: String
+input = "...........\n.....###.#.\n.###.##..#.\n..#.#...#..\n....#.#....\n.##..S####.\n.##..#...#.\n.......##..\n.##.#.####.\n.##..##.##.\n..........."
+
+test1 :: Expectation
+test1 = part1 input `shouldBe` "42"
+
+test2 :: Expectation
+test2 = part2 input `shouldBe` "528192899606863"
+
+test :: IO ()
+test = hspec $ do
+    describe "day21" $ do
+        describe "part1" $ do
+            it "should work for the examples" test1
+
+        describe "part2" $ do
+            it "should work for the examples" test2