2024 Day 05 in Haskell

2024/day05/day05.hs

@@ -0,0 +1,45 @@
+import Aoc
+import Data.List.Split ( splitOn )
+import Control.Arrow ( second, first )
+import Data.List ( nub, partition, sort )
+import Data.Maybe ( fromMaybe )
+
+parseFile :: [String] -> ([(Int, [Int])], [[Int]])
+parseFile lines =
+    let 
+        [rules, pages] = splitOn [""] lines
+        rules' = map (pair . map read . splitOn "|") rules
+        mentioned = nub $ concatMap unpair rules'
+        getAll a = map snd $ filter ((==a) . fst) rules' 
+    in
+    (zip mentioned (map getAll mentioned), map (map read . splitOn ",") pages)
+    
+takeMiddle :: [a] -> a
+takeMiddle xs = xs !! (length xs `div` 2)
+
+isSafe :: [(Int, [Int])] -> [Int] -> Bool
+isSafe rules = fst . foldl isSafe' (True, [])
+    where
+        isSafe' (False, prev) _ = (False, prev)
+        isSafe' (True, prev) a = (fromMaybe True $ not . any (`elem` prev) <$> lookup a rules, a:prev)
+
+part1 :: ([(Int, [Int])], [[Int]]) -> (Int, [[Int]])
+part1 (rules, pages) = first (sum . map takeMiddle) $ partition (isSafe rules) pages
+
+part2 :: ([(Int, [Int])], [[Int]]) -> [[Int]] -> Int
+part2 (rules, _) = sum . map (takeMiddle . reorder)
+    where
+        getPosM rules = memoise $ \values n ->
+            let ns = fromMaybe [] $ lookup n rules in
+            if null ns 
+                then (0, values)
+                else
+                    let (ns', values') = memoMap (getPosM rules) values ns in
+                    (maximum ns' + 1, values')
+
+        reorder ns =
+            let rules' = map (second $ filter (`elem` ns)) $ filter ((`elem` ns) . fst) rules in
+            map snd $ sort $ zip (unmemo (memoMap $ getPosM rules') ns) ns
+
+main :: IO ()
+main = aocMain' parseFile part1 part2

aoc.hs

@@ -60,6 +60,8 @@ starsAndBars 0 b = [replicate (b+1) 0]
 starsAndBars n 0 = [[n]]
 starsAndBars n b = map (0:) (starsAndBars n (b-1)) ++ map (\xs -> head xs + 1 : tail xs) (starsAndBars (n-1) b)
 
+-- Kernel
+
 kernelMapRow' :: ([[a]] -> b) -> Maybe a -> [a] -> [a] -> [a] -> [b]
 kernelMapRow' kernel x
     (a:r0@(b:c:_))
@@ -81,3 +83,23 @@ kernelMap' kernel Nothing [_, _] = []
 kernelMap :: ([[a]] -> b) -> Maybe a -> [[a]] -> [[b]]
 kernelMap kernel (Just x) as = kernelMap' kernel (Just x) ((map (const x) (head as)):as)
 kernelMap kernel Nothing as = kernelMap' kernel Nothing as
+
+-- Memoisation
+
+type MemoF' a b a' b' = Eq a' => [(a', b')] -> a -> (b, [(a', b')])
+type MemoF a b = MemoF' a b a b
+
+memoise :: Eq a => MemoF a b -> MemoF a b
+memoise f values a = case lookup a values of
+    Just b -> (b, values)
+    Nothing -> let (b, values') = f values a in (b, (a, b):values')
+
+memoMap :: MemoF' a b a' b' -> MemoF' [a] [b] a' b'
+memoMap f values [] = ([], values)
+memoMap f values (a:as) =
+    let (b, values') = f values a in
+    let (bs, values'') = memoMap f values' as in
+    (b:bs, values'')
+
+unmemo :: Eq a' => MemoF' a b a' b' -> a -> b
+unmemo f = fst . f []