faster

glguy · Dec 20, 2024 · 60390c1 · 60390c1
1 parent d792b43
commit 60390c1
Showing 1 changed file with 11 additions and 5 deletions.
diff --git a/solutions/src/2024/20.hs b/solutions/src/2024/20.hs
@@ -11,9 +11,9 @@ Maintainer  : emertens@gmail.com
 module Main (main) where
 
 import Advent (getInputArray, arrIx, count)
-import Advent.Coord (Coord, cardinal, manhattan)
+import Advent.Coord (Coord(C), cardinal, manhattan, coordRow, coordCol)
 import Advent.Search (dfs)
-import Data.Array.Unboxed (UArray, amap, assocs)
+import Data.Array.Unboxed (UArray, amap, assocs, accumArray, bounds)
 import Data.List (tails)
 
 -- >>> :main
@@ -26,10 +26,16 @@ main =
         start : _ = [p | (p, 'S') <- assocs input]
         step p    = [p' | p' <- cardinal p, True <- arrIx open p']
         path      = dfs step start
+        pathArray = accumArray (\_ e -> e) (-1) (bounds open) (zip path [0..]) :: UArray Coord Int
+        (C loy lox, C hiy hix) = bounds input
         cheats    = [ d
-                    | (p1, c1) : more <- tails (zip path [0..])
-                    , (p2, c2)        <- drop 100 more
-                    , let d = manhattan p1 p2, d <= 20
+                    | (p1, c1) <- zip path [0..]
+                    , y2 <- [max loy (coordRow p1 - 20) .. min hiy (coordRow p1 + 20)]
+                    , let dy = abs (coordRow p1 - y2)
+                    , x2 <- [max lox (coordCol p1 - 20 + dy) .. min hix (coordCol p1 + 20 - dy)]
+                    , let dx = abs (coordCol p1 - x2)
+                    , let d = dx + dy
+                    , c2 <- arrIx pathArray (C y2 x2)
                     , c2 - c1 >= 100 + d
                     ]
     print (count 2 cheats)