Merge pull request #41 from ephemient/hs/day5

hs/aoc2024.cabal

@@ -35,6 +35,7 @@ library
     containers ^>=0.7,
     megaparsec ^>=9.7.0,
     text ^>=2.1.2,
+    vector ^>=0.13.2.0,
 
   ghc-options: -Wall
   default-language: GHC2024

hs/src/Day5.hs

@@ -3,56 +3,59 @@
 -- Description:    <https://adventofcode.com/2024/day/5 Day 5: Print Queue>
 module Day5 (part1, part2) where
 
-import Control.Arrow (second)
-import Data.IntMap qualified as IntMap (findWithDefault, fromList, (!?))
-import Data.IntMap.Strict qualified as IntMap (fromListWith)
-import Data.IntSet qualified as IntSet (empty, member, singleton, toList, union)
-import Data.List (sortBy)
-import Data.Maybe (fromJust, fromMaybe)
+import Data.List (tails)
+import Data.Set (Set)
+import Data.Set qualified as Set (fromList, notMember)
 import Data.String (IsString)
 import Data.Text (Text)
+import Data.Vector.Generic (Vector)
+import Data.Vector.Generic qualified as V (fromList, length, modify, (!))
+import Data.Vector.Generic.Mutable qualified as MV (length, read, write)
+import Data.Vector.Unboxed qualified as UV (Vector)
 import Data.Void (Void)
 import Text.Megaparsec (MonadParsec, ParseErrorBundle, Stream (Token, Tokens), parse, sepEndBy, sepEndBy1, skipMany)
 import Text.Megaparsec.Char (char, newline)
 import Text.Megaparsec.Char.Lexer (decimal)
 
-parser :: (MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a) => m ([(a, a)], [[a]])
+parser :: (MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a, Ord a) => m (Set (a, a), [[a]])
 parser =
-  (,)
+  (,) . Set.fromList
     <$> ((,) <$> decimal <* char '|' <*> decimal) `sepEndBy` newline
     <* skipMany newline
     <*> (decimal `sepEndBy1` char ',') `sepEndBy` newline
 
 part1 :: Text -> Either (ParseErrorBundle Text Void) Int
 part1 input = do
   (deps, updates) <- parse parser "" input
-  pure $
-    sum
-      [ update !! (length update `div` 2)
-      | update <- updates,
-        let order = IntMap.fromList @Int $ zip update [0 ..],
-        and
-          [ fromMaybe True $ (<) <$> order IntMap.!? a <*> order IntMap.!? b
-          | (a, b) <- deps
-          ]
-      ]
+  pure . sum $
+    [ update !! (length update `div` 2)
+    | update <- updates,
+      and [(b, a) `Set.notMember` deps | a : rest <- tails update, b <- rest]
+    ]
 
 part2 :: Text -> Either (ParseErrorBundle Text Void) Int
 part2 input = do
   (deps, updates) <- parse parser "" input
-  let deps' = IntMap.fromListWith IntSet.union $ second IntSet.singleton <$> deps
-      tryCompare a b
-        | a == b = Just EQ
-        | b `IntSet.member` (IntMap.findWithDefault IntSet.empty a deps') = Just LT
-        | a `IntSet.member` (IntMap.findWithDefault IntSet.empty b deps') = Just GT
-        | otherwise =
-            mconcat (tryCompare a <$> maybe [] IntSet.toList (deps' IntMap.!? b))
-              <> fmap (compare EQ) (mconcat (tryCompare b <$> maybe [] IntSet.toList (deps' IntMap.!? a)))
-      compare' a b = fromJust $ tryCompare a b
-  pure $
-    sum
-      [ update' !! (length update `div` 2)
-      | update <- updates,
-        let update' = sortBy compare' update,
-        update /= update'
-      ]
+  let ok a b = (b, a) `Set.notMember` deps
+  pure . sum $
+    [ pages V.! (V.length pages `div` 2)
+    | update <- V.fromList @UV.Vector @Int <$> updates,
+      let pages = sort' ok update,
+      update /= pages
+    ]
+
+sort' :: (Vector v a) => (a -> a -> Bool) -> v a -> v a
+sort' ok = V.modify $ \v ->
+  let go i j
+        | j < MV.length v = do
+            x <- MV.read v i
+            y <- MV.read v j
+            if ok x y
+              then go i (j + 1)
+              else do
+                MV.write v i y
+                MV.write v j x
+                go i (i + 1)
+        | i < MV.length v = go (i + 1) (i + 2)
+        | otherwise = pure ()
+   in go 0 1