clean up day 17 a bit

mstksg · Dec 17, 2023 · eb6783b · eb6783b
1 parent bc42459
commit eb6783b
Showing 1 changed file with 47 additions and 67 deletions.
diff --git a/src/AOC/Challenge/Day17.hs b/src/AOC/Challenge/Day17.hs
@@ -1,4 +1,4 @@
-{-# OPTIONS_GHC -Wno-unused-imports   #-}
+{-# OPTIONS_GHC -Wno-unused-imports #-}
 {-# OPTIONS_GHC -Wno-unused-top-binds #-}
 
 -- |
@@ -20,76 +20,56 @@
 --     types @_ :~> _@ with the actual types of inputs and outputs of the
 --     solution.  You can delete the type signatures completely and GHC
 --     will recommend what should go in place of the underscores.
+module AOC.Challenge.Day17
+  ( day17a,
+    day17b,
+  )
+where
 
-module AOC.Challenge.Day17 (
-    day17a
-  , day17b
-  ) where
-
-import           AOC.Prelude
-
-import qualified Data.Graph.Inductive           as G
-import qualified Data.IntMap                    as IM
-import qualified Data.IntSet                    as IS
-import qualified Data.List.NonEmpty             as NE
-import qualified Data.List.PointedList          as PL
+import AOC.Prelude
+import qualified Data.Graph.Inductive as G
+import qualified Data.IntMap as IM
+import qualified Data.IntSet as IS
+import qualified Data.List.NonEmpty as NE
+import qualified Data.List.PointedList as PL
 import qualified Data.List.PointedList.Circular as PLC
-import qualified Data.Map                       as M
-import qualified Data.OrdPSQ                    as PSQ
-import qualified Data.Sequence                  as Seq
-import qualified Data.Set                       as S
-import qualified Data.Text                      as T
-import qualified Data.Vector                    as V
-import qualified Linear                         as L
-import qualified Text.Megaparsec                as P
-import qualified Text.Megaparsec.Char           as P
-import qualified Text.Megaparsec.Char.Lexer     as PP
+import qualified Data.Map as M
+import qualified Data.OrdPSQ as PSQ
+import qualified Data.Sequence as Seq
+import qualified Data.Set as S
+import qualified Data.Text as T
+import qualified Data.Vector as V
+import qualified Linear as L
+import qualified Text.Megaparsec as P
+import qualified Text.Megaparsec.Char as P
+import qualified Text.Megaparsec.Char.Lexer as PP
 
-day17a :: _ :~> _
-day17a = MkSol
-    { sParse = noFail $ parseAsciiMap digitToIntSafe
-    , sShow  = show
-    , sSolve = \mp -> do
+day17 :: Int -> Int -> Map Point Int :~> Int
+day17 a b =
+  MkSol
+    { sParse = noFail $ parseAsciiMap digitToIntSafe,
+      sShow = show,
+      sSolve = \mp -> do
         V2 p1 p2 <- boundingBox' (M.keys mp)
-        let grow (lastThree, p) = M.fromList
-              [ ((take 3 (d : lastThree), p'), cost)
-                | d <- toList allDir
-              , case lastThree of
-                  [] -> True
-                  dd : _ -> dd /= d <> South
-              , lastThree /= replicate 3 d
-              , let p' = p + dirPoint d
-              , cost <- maybeToList $ M.lookup p' mp
-              ]
-        fst <$> aStar ((`mannDist` p2) . snd) grow ([],p1) ((== p2) . snd)
+        let grow (lastDir, p) =
+              M.fromList
+                [ ((Just d, p'), sum (snd <$> steps))
+                  | d <- toList allDir,
+                    case lastDir of
+                      Nothing -> True
+                      Just d' -> d /= d' && d /= (d' <> South),
+                    steps <- drop a $ inits do
+                      i <- [1 .. b]
+                      let p' = p + i *^ dirPoint d
+                      cost <- maybeToList $ M.lookup p' mp
+                      pure (p', cost),
+                    (p', _) <- maybeToList $ lastMay steps
+                ]
+        fst <$> aStar ((`mannDist` p2) . snd) grow (Nothing, p1) ((== p2) . snd)
     }
 
--- aStar
---     :: forall n p. (Ord n, Ord p, Num p)
---     => (n -> p)         -- ^ heuristic
---     -> (n -> Map n p)   -- ^ neighborhood
---     -> n                -- ^ start
---     -> (n -> Bool)      -- ^ target
---     -> Maybe (p, [n])   -- ^ the shortest path, if it exists, and its cost
+day17a :: Map Point Int :~> Int
+day17a = day17 1 3
 
-day17b :: _ :~> _
-day17b = MkSol
-    { sParse = sParse day17a
-    , sShow  = show
-    , sSolve = \mp -> do
-        V2 p1 p2 <- boundingBox' (M.keys mp)
-        let grow (lastDir, p) = M.fromList
-              [ ((Just d, p'), sum (snd <$> steps))
-                | d <- toList allDir
-              , case lastDir of
-                  Nothing -> True
-                  Just d' -> d /= d' && d /= (d' <> South)
-              , steps <- drop 4 $ inits do
-                  i <- [1..10]
-                  let p' = p + i *^ dirPoint d
-                  cost <- maybeToList $ M.lookup p' mp
-                  pure (p', cost)
-              , (p', _) <- maybeToList $ lastMay steps
-              ]
-        fst <$> aStar ((`mannDist` p2) . snd) grow (Nothing,p1) ((== p2) . snd)
-    }
+day17b :: Map Point Int :~> Int
+day17b = day17 4 10