aoc-2024-haskell/day5/Main.hs

module Main where

import System.Environment (getArgs)
import Data.Maybe (mapMaybe)
import Data.List (find)

import Text.Parsec (parse, Parsec, endBy, sepBy1, char, newline, many)
import Text.ParserCombinators.Parsec.Number (decimal)

type OrderingRule = (Int, Int)
type PagesUpdate = [Int]

printInfo :: Parsec String () ([OrderingRule], [PagesUpdate])
printInfo = (,) <$> orderingRules <*> (newline *> pagesUpdates) <* many (char '\n')
    where
        pagesUpdates :: Parsec String () [PagesUpdate]
        pagesUpdates = endBy (sepBy1 decimal $ char ',') newline

        orderingRules :: Parsec String () [OrderingRule]
        orderingRules = endBy ((,) <$> decimal <*> (char '|' *> decimal)) $ char '\n'

allDependencies :: [a] -> [(a, a)]
allDependencies (h:tl) = map (flip (,) h) tl ++ allDependencies tl
allDependencies [] = []

middleElement :: [a] -> Maybe a
middleElement [] = Nothing
middleElement [a] = Just a
middleElement (_:h':tl) = case reverse (h':tl) of
    (_:middle) -> middleElement middle
    [] -> Nothing

isWellOrdered :: [OrderingRule] -> PagesUpdate -> Bool
isWellOrdered orderingRules = all (`notElem` orderingRules) . allDependencies

part1 :: [OrderingRule] -> [PagesUpdate] -> Int
part1 orderingRules pagesUpdates = (sum . mapMaybe middleElement) onlyWellOrdered
    where onlyWellOrdered = filter (isWellOrdered orderingRules) pagesUpdates

partitions :: [a] -> [([a], a, [a])]
partitions (h:tl) = ([], h, tl) : map (\(as, a, bs) -> (h:as, a, bs)) (partitions tl)
partitions [] = []

doublePartitions :: [a] -> [([a], a, [a], a, [a])]
doublePartitions ls = [(as, a, bs, b, cs)
    | (as, a, bs') <- partitions ls
    , (bs, b, cs) <- partitions bs']

reorder :: [OrderingRule] -> PagesUpdate -> PagesUpdate
reorder orderingRules pagesUpdate = case find badlyOrdered $ doublePartitions pagesUpdate of
    Just (as, a, bs, b, cs) -> reorder orderingRules (as ++ b : bs ++ a : cs)
    Nothing -> pagesUpdate
    where
        badlyOrdered :: ([Int], Int, [Int], Int, [Int]) -> Bool
        badlyOrdered (_, a, _, b, _) = (b, a) `elem` orderingRules

part2 :: [OrderingRule] -> [PagesUpdate] -> Int
part2 orderingRules = sum
    . mapMaybe (middleElement . reorder orderingRules)
    . filter (not . isWellOrdered orderingRules)

main :: IO ()
main = do
    args <- getArgs
    (fp:_) <- return args
    content <- readFile fp
    input <- either (fail . show) return $ parse printInfo fp content
    -- putStrLn $ "Solution (Part 1) : " ++ show (uncurry part1 input)
    putStrLn $ "Solution (Part 2) : " ++ show (uncurry part2 input)