Merge pull request #177 from ephemient/hs/day6

ephemient · Dec 27, 2024 · 0de1636 · 0de1636
2 parents 16f9832 + c2a783c
commit 0de1636
Showing 1 changed file with 19 additions and 14 deletions.
diff --git a/hs/src/Day6.hs b/hs/src/Day6.hs
@@ -4,18 +4,17 @@
 module Day6 (part1, part2) where
 
 import Control.Monad (ap)
-import Control.Parallel.Strategies (parList, rseq, withStrategy)
+import Control.Parallel.Strategies (parMap, rseq)
 import Data.Containers.ListUtils (nubOrd)
-import Data.List ((\\))
 import Data.List.NonEmpty (NonEmpty ((:|)))
 import Data.Maybe (catMaybes, isJust)
 import Data.Semigroup (Max (Max), sconcat)
 import Data.Set (Set)
-import Data.Set qualified as Set (empty, insert, member, singleton)
+import Data.Set qualified as Set (empty, insert, lookupGE, member, notMember, singleton)
 import Data.Text (Text)
 import Data.Text qualified as T (lines, unpack)
 
-parse :: Text -> ((Int, Int), Set (Int, Int), [(Int, Int)])
+parse :: Text -> ((Int, Int), Set (Int, Int), [((Int, Int), (Int, Int))])
 parse input = ((maxY, maxX), blocks, start)
   where
     (Max maxY, Max maxX, (blocks, start)) =
@@ -24,20 +23,23 @@ parse input = ((maxY, maxX), blocks, start)
           :| [ ( Max y,
                  Max x,
                  case char of
-                   '^' -> (mempty, [(y, x)])
+                   '^' -> (mempty, [((y, x), (-1, 0))])
+                   '<' -> (mempty, [((y, x), (0, -1))])
+                   '>' -> (mempty, [((y, x), (0, 1))])
+                   'v' -> (mempty, [((y, x), (1, 0))])
                    '#' -> (Set.singleton (y, x), mempty)
                    _ -> mempty
                )
              | (y, line) <- zip [0 ..] $ T.lines input,
                (x, char) <- zip [0 ..] $ T.unpack line
              ]
 
-visited :: (Int, Int) -> Set (Int, Int) -> (Int, Int) -> [((Int, Int), (Int, Int))]
-visited (maxY, maxX) blocks start = catMaybes $ takeWhile isJust $ iterate (>>= step) $ Just (start, (-1, 0))
+visited :: (Int, Int) -> Set (Int, Int) -> ((Int, Int), (Int, Int)) -> [((Int, Int), (Int, Int))]
+visited (maxY, maxX) blocks start = catMaybes $ takeWhile isJust $ iterate (>>= step) $ Just start
   where
     step (pos@(y, x), d@(dy, dx))
       | y' < 0 || maxY < y' || x' < 0 || maxX < x' = Nothing
-      | (y', x') `Set.member` blocks = step (pos, (dx, -dy))
+      | (y', x') `Set.member` blocks = Just (pos, (dx, -dy))
       | otherwise = Just ((y', x'), d)
       where
         y' = y + dy
@@ -50,11 +52,14 @@ part1 input = length $ nubOrd $ map fst $ start >>= visited maxes blocks
 
 part2 :: Text -> Int
 part2 input =
-  length . filter id . withStrategy (parList rseq) $
-    [ isLoop (Set.insert add blocks) pos0
-    | pos0 <- start,
-      add <- nubOrd (fst <$> visited maxes blocks pos0) \\ [pos0]
-    ]
+  length . filter id . parMap rseq isLoop $
+    start >>= (zip `ap` scanl (flip Set.insert) Set.empty) . visited maxes blocks
   where
     (maxes, blocks, start) = parse input
-    isLoop blocks' pos0 = any (uncurry Set.member) $ zip `ap` scanl (flip Set.insert) Set.empty $ visited maxes blocks' pos0
+    isLoop (start'@((y, x), (dy, dx)), seen) =
+      pos' `Set.notMember` blocks
+        && Just pos' /= (fst <$> Set.lookupGE (pos', minBound) seen)
+        && or (zipWith Set.member `ap` scanl (flip Set.insert) seen $ visited maxes blocks' start')
+      where
+        pos' = (y + dy, x + dx)
+        blocks' = Set.insert pos' blocks