cleanup

solutions/src/2024/20.hs

@@ -14,29 +14,31 @@ module Main (main) where
 import Advent (getInputArray, arrIx)
 import Advent.Coord (cardinal, Coord, manhattan)
 import Advent.Search (astar, AStep(AStep))
-import Data.Array.Unboxed (UArray, amap, assocs)
+import Data.Array.Unboxed (UArray, amap, assocs, accumArray, bounds, (!))
+import Data.List (tails)
 import Data.Map qualified as Map
 
+-- >>> :main
+-- 1346
+-- 985482
 main :: IO ()
 main =
  do input <- getInputArray 2024 20
-    let walls = amap ('#' /=) input
-    let start : _ = [p | (p,'S') <- assocs input]
-    let end   : _ = [p | (p,'E') <- assocs input]
-    let optimals = Map.fromList (search walls end)
-    let part n = length [()
-                | (p1, True) <- assocs walls
-                , (p2, True) <- assocs walls
-                , let d = manhattan p1 p2
-                , d > 0, d <= n
-                , let c1 = optimals Map.! p1
-                , let c2 = optimals Map.! p2
-                , c1 - c2 - d >= 100
-                ]
-    print (part  2)
-    print (part 20)
+    let open      = amap ('#' /=) input
+        start : _ = [p | (p,'S') <- assocs input]
+        end   : _ = [p | (p,'E') <- assocs input]
+        optimals  = search open end
+        cheat n   = length [()
+                  | (p1,c1) : more <- tails optimals
+                  , (p2,c2)        <- more
+                  , let d = manhattan p1 p2
+                  , d >= 2, d <= n
+                  , abs (c1 - c2) >= 100 + d
+                  ]
+    print (cheat  2)
+    print (cheat 20)
 
 search :: UArray Coord Bool -> Coord -> [(Coord, Int)]
-search input start = astar step start
+search open = astar step
   where
-    step p = [ AStep p' 1 0 | p' <- cardinal p, True <- arrIx input p']
+    step p = [AStep p' 1 0 | p' <- cardinal p, True <- arrIx open p']