Day 14: Look for combination of max per-axis safety

Author: ephemient

hs/aoc2024.cabal

@@ -21,6 +21,7 @@ source-repository head
 library
   hs-source-dirs: src
   exposed-modules:
+    Common
     Day1
     Day10
     Day11
@@ -47,9 +48,6 @@ library
     Day8
     Day9
 
-  other-modules:
-    Common
-
   build-depends:
     array ^>=0.5.7.0,
     async ^>=2.2.5,
@@ -90,6 +88,7 @@ test-suite aoc2024-test
   hs-source-dirs: test
   main-is: Main.hs
   other-modules:
+    CommonSpec
     Day10Spec
     Day11Spec
     Day12Spec
@@ -117,6 +116,7 @@ test-suite aoc2024-test
     Day9Spec
 
   build-depends:
+    QuickCheck ^>=2.15.0.1,
     aoc2024,
     base >=4.20 && <4.22,
     hspec ^>=2.11.10,

hs/src/Common.hs

@@ -1,18 +1,38 @@
-module Common (groupConsecutiveBy, readEntire, readMany, readSome) where
+module Common (crt, egcd, readEntire, readMany, readSome) where
 
-import Control.Arrow (first)
+import Data.Bifunctor (first)
 import Data.Char (isSpace)
 import Data.List.NonEmpty (NonEmpty ((:|)))
 import Data.Text (Text)
 import Data.Text qualified as T (dropWhile, null)
 import Data.Text.Read (Reader)
 
-groupConsecutiveBy :: (a -> a -> Bool) -> [a] -> [[a]]
-groupConsecutiveBy f xs = chunk id $ zip xs $ True : zipWith f xs (drop 1 xs)
+-- | Chinese remainder theorem.
+--
+--  prop> crt (r1, q1) (r2, q2) == (r3, q3) ==>
+--            r3 `mod` q1 == r1 && q3 `mod` q1 == 0 &&
+--            r3 `mod` q2 == r2 && q3 `mod` q2 == 0
+crt :: (Integral a) => (a, a) -> (a, a) -> (a, a)
+crt (r1, q1) (r2, q2) = (r3 `mod` q3, q3)
   where
-    chunk k [] = filter (not . null) [k []]
-    chunk k ((x, False) : rest) = k [] : chunk (x :) rest
-    chunk k ((x, True) : rest) = chunk (k . (x :)) rest
+    q3 = lcm q1 q2
+    -- r3 * q2 == r1 * q2 (mod q3)
+    -- r3 * q3 == r2 * q1 (mod q3)
+    -- r3 * (q1 + q2) = r1 * q2 + r2 * q1 (mod q3)
+    (t, _, g) = egcd (q1 + q2) q3
+    -- t * (q1 + q2) == g (mod q3)
+    -- r3 = (r1 * q2 + r2 * q1) * t / g (mod q3)
+    (r3, 0) = ((r1 * q2 + r2 * q1) * t) `divMod` g
+
+-- | Extended GCD.
+--
+--  prop> gcd a b == (s, t, g) ==> a * s + b * t == g
+egcd :: (Integral a) => a -> a -> (a, a, a)
+egcd a 0 = (1, 0, a)
+egcd a b = (t, s - q * t, g)
+  where
+    (q, r) = a `quotRem` b
+    (s, t, g) = egcd b r
 
 readEntire :: Reader a -> Text -> Either String a
 readEntire reader input = do

hs/src/Day14.hs

@@ -5,72 +5,58 @@
 -- Description:    <https://adventofcode.com/2024/day/14 Day 14: Restroom Redoubt>
 module Day14 (part1, part1', part2) where
 
-import Common (groupConsecutiveBy)
-import Control.Monad (join, liftM2)
-import Control.Parallel.Strategies (parList, rdeepseq, withStrategy)
-import Data.Char (intToDigit)
-import Data.Map qualified as Map (findWithDefault)
-import Data.Map.Strict qualified as Map (fromListWith)
-import Data.Ord (Down (Down))
-import Data.Set qualified as Set (fromList, toList)
+import Common (crt)
+import Control.Monad (ap)
+import Data.Foldable (foldMap')
+import Data.Semigroup (Arg (Arg), Sum (Sum))
 import Data.String (IsString)
 import Data.Text (Text)
 import Data.Void (Void)
-import Debug.Trace (traceM)
 import Text.Megaparsec (MonadParsec, ParseErrorBundle, Stream (Token, Tokens), parse, sepEndBy1)
 import Text.Megaparsec.Char (char, newline, string)
 import Text.Megaparsec.Char.Lexer qualified as L (decimal, signed)
 
 parser :: (MonadParsec e s m, IsString (Tokens s), Token s ~ Char, Num a) => m [((a, a), (a, a))]
 parser = line `sepEndBy1` newline
   where
-    line = (,) <$> (string "p=" *> v2) <*> (string " v=" *> v2)
-    v2 = (,) <$> (L.signed (pure ()) L.decimal <* char ',') <*> L.signed (pure ()) L.decimal
+    line = do
+      x <- string "p=" *> L.signed (pure ()) L.decimal
+      y <- char ',' *> L.signed (pure ()) L.decimal
+      vx <- string " v=" *> L.signed (pure ()) L.decimal
+      vy <- char ',' *> L.signed (pure ()) L.decimal
+      pure ((x, vx), (y, vy))
 
 part1 :: Text -> Either (ParseErrorBundle Text Void) Int
 part1 = part1' 101 103
 
 part1' :: Int -> Int -> Text -> Either (ParseErrorBundle Text Void) Int
 part1' width height input = do
   robots <- parse parser "" input
-  let totals =
-        Map.fromListWith (+) $
-          [ ((compare x $ width `div` 2, compare y $ height `div` 2), 1)
-          | ((x0, y0), (vx, vy)) <- robots,
-            let x = (x0 + vx * t) `mod` width
-                y = (y0 + vy * t) `mod` height
-          ]
-  pure $ product [Map.findWithDefault 0 k totals | k <- join (liftM2 (,)) [LT, GT]]
+  let (Sum q1, Sum q2, Sum q3, Sum q4) = foldMap' f robots
+  pure $ q1 * q2 * q3 * q4
   where
     t = 100
+    f ((x, vx), (y, vy)) = case ( compare ((x + t * vx) `mod` width) (width `div` 2),
+                                  compare ((y + t * vy) `mod` height) (height `div` 2)
+                                ) of
+      (LT, LT) -> (Sum 1, Sum 0, Sum 0, Sum 0)
+      (LT, GT) -> (Sum 0, Sum 1, Sum 0, Sum 0)
+      (GT, LT) -> (Sum 0, Sum 0, Sum 1, Sum 0)
+      (GT, GT) -> (Sum 0, Sum 0, Sum 0, Sum 1)
+      _ -> mempty
 
 part2 :: Text -> Either (ParseErrorBundle Text Void) Int
 part2 input = do
-  robots <- parse parser "" input
-  let (_, bestTime) =
-        minimum . withStrategy (parList rdeepseq) $
-          [ (Down $ maximum $ map length verticalLines, t)
-          | t <- [0 .. lcm width height - 1],
-            let verticalLines =
-                  groupConsecutiveBy isLine . Set.toList . Set.fromList $
-                    [ ((y0 + vy * t) `mod` height, (x0 + vx * t) `mod` width)
-                    | ((x0, y0), (vx, vy)) <- robots
-                    ]
-                isLine (y0, x0) (y1, x1) = y0 == y1 && x0 + 1 == x1
-          ]
-      positions =
-        Map.fromListWith (+) $
-          [ (((x0 + vx * bestTime) `mod` width, (y0 + vy * bestTime) `mod` height), 1)
-          | ((x0, y0), (vx, vy)) <- robots
-          ]
-      line y =
-        [ case Map.findWithDefault 0 (x, y) positions of
-            0 -> '.'
-            n -> if n < 10 then intToDigit n else '+'
-        | x <- [0 .. width - 1]
-        ]
-  mapM_ (traceM . line) [0 .. height - 1]
-  pure bestTime
+  (xrobots, yrobots) <- unzip <$> parse parser "" input
+  let Arg _ x = maximum $ (flip Arg `ap` score xrobots width) <$> [0 .. width - 1]
+      Arg _ y = maximum $ (flip Arg `ap` score yrobots height) <$> [0 .. height - 1]
+  pure $ fst $ crt (x, width) (y, height)
   where
-    width = 101
-    height = 103
+    (width, height) = (101, 103)
+    score robots m t = max h1 h2 :: Int
+      where
+        (Sum h1, Sum h2) = foldMap' f robots
+        f (p, v) = case compare ((p + t * v) `mod` m) (m `div` 2) of
+          LT -> (Sum 1, Sum 0)
+          EQ -> (Sum 0, Sum 0)
+          GT -> (Sum 0, Sum 1)

hs/test/CommonSpec.hs

@@ -0,0 +1,36 @@
+module CommonSpec (spec) where
+
+import Common (crt, egcd)
+import Test.Hspec (Spec, describe)
+import Test.Hspec.QuickCheck (prop)
+import Test.QuickCheck (Large (Large), arbitrarySizedIntegral, conjoin, counterexample, cover, property, suchThat, (===))
+
+spec :: Spec
+spec = do
+  describe "crt" $ do
+    prop "is correct" $ do
+      q1 <- arbitrarySizedIntegral `suchThat` (/= 0)
+      q2 <- arbitrarySizedIntegral `suchThat` (/= 0)
+      let q = gcd q1 q2
+      r1 <- arbitrarySizedIntegral `suchThat` ((== 0) . (`mod` q))
+      r2 <- arbitrarySizedIntegral `suchThat` ((== 0) . (`mod` q))
+      let (r3, q3) = crt @Int (r1, q1) (r2, q2)
+      pure
+        . cover 90 (abs q1 /= 1 && abs q1 /= 1) "non-trivial"
+        . counterexample ("(r1,q1) = " ++ show (r1, q1))
+        . counterexample ("(r2,q2) = " ++ show (r2, q2))
+        . counterexample ("(r3,q3) = " ++ show (r3, q3))
+        $ conjoin
+          [ counterexample "r1 == r3 (mod q1)" $ r1 `mod` q1 === r3 `mod` q1,
+            counterexample "q3 == 0 (mod q1)" $ q3 `mod` q1 === 0,
+            counterexample "r2 == r3 (mod q2)" $ r2 `mod` q2 === r3 `mod` q2,
+            counterexample "q3 == 0 (mod q2)" $ q3 `mod` q2 === 0,
+            counterexample "|r3| < |q3|" $ property $ abs r3 < abs q3,
+            counterexample "|gcd(r1 - r2, q1, q2)| == |q1 * q2 / q3|" $
+              abs (gcd (r1 - r2) $ gcd q1 q2) === abs (q1 * q2 `div` q3)
+          ]
+  describe "egcd" $ do
+    prop "is correct" $ \(Large a) (Large b) ->
+      cover 90 (abs a > 1 && abs b > 1) "non-trivial" $
+        let (s, t, g) = egcd @Int a b
+         in a * s + b * t === g