advent25/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2024/01/06/advent-of-code-2023-day-25/
   2
   3 import AoC
   4
   5 import Data.Text (Text)
   6 import qualified Data.Text.IO as TIO
   7 import Data.Attoparsec.Text hiding (take)
   8 -- import Control.Applicative
   9 import Data.List (foldl', unfoldr, sort, delete)
  10 import qualified Data.Map.Strict as M
  11 import Data.Map.Strict ((!))
  12 import qualified Data.Set as S
  13 import Data.Set ((\\))
  14 import qualified Data.Sequence as Q
  15 import Data.Sequence ( (><), Seq(..) )
  16 import Data.Foldable (toList)
  17 import Data.Maybe (catMaybes)
  18 import System.Random
  19 import qualified Data.MultiSet as MS
  20 import Data.Tuple (swap)
  21
  22 type Graph = M.Map String [String]
  23
  24 main :: IO ()
  25 main =
  26   do  dataFileName <- getDataFileName
  27       text <- TIO.readFile dataFileName
  28       let halfGraph = successfulParse text
  29       let graph = reverseGraph halfGraph
  30       print $ part1 graph
  31
  32 part1 :: Graph -> Int
  33 part1 graph = (S.size componentA) * (S.size componentB)
  34   where (ss, gs) = splitAt 200 $ randomNodes graph 400
  35         paths = fmap toList $ catMaybes $ fmap (bfsPair graph) $ zip ss gs
  36         pathCounts = MS.fromList $ concatMap getEdges paths
  37         populars = fmap snd $ take 3 $ reverse $ sort $ fmap swap $ MS.toOccurList pathCounts
  38         separatedGraph = foldl' removeEdge graph populars
  39         (a, b) = head populars
  40         componentA = componentOf separatedGraph (S.singleton a) S.empty
  41         componentB = componentOf separatedGraph (S.singleton b) S.empty
  42
  43
  44 reverseGraph :: Graph -> Graph
  45 reverseGraph graph = M.foldlWithKey' reverseNode graph graph
  46
  47 reverseNode :: Graph -> String -> [String] -> Graph
  48 reverseNode graph here theres = foldl' (addReversed here) graph theres
  49
  50 addReversed :: String -> Graph -> String -> Graph
  51 addReversed here graph there = M.insertWith (++) there [here] graph
  52
  53 bfsPair :: Graph -> (String, String) -> Maybe (Seq String)
  54 bfsPair graph (start, goal) = bfs graph (Q.singleton (Q.singleton start)) goal S.empty
  55
  56 bfs :: Graph -> Seq (Seq String) -> String -> S.Set String -> Maybe (Seq String)
  57 bfs _ Q.Empty _ _ = Nothing
  58 bfs graph (current :<| agenda) goal closed
  59   | here == goal = Just current
  60   | here `S.member` closed = bfs graph agenda goal closed
  61   | otherwise = bfs graph (agenda >< nexts) goal (S.insert here closed)
  62   where (_ :|> here) = current
  63         nexts = Q.fromList $ fmap (current :|>) $ graph ! here
  64
  65 componentOf :: Graph -> S.Set String -> S.Set String -> S.Set String
  66 componentOf graph boundary0 found
  67   | S.null boundary0 = found
  68   | otherwise = componentOf graph boundary2 found'
  69   where (here, boundary1) = S.deleteFindMin boundary0
  70         found' = S.insert here found
  71         boundary2 = S.union boundary1 $ (S.fromList $ graph ! here) \\ found'
  72
  73 randomNodes :: Graph -> Int -> [String]
  74 randomNodes graph n = fmap (\i -> fst $ M.elemAt i graph) indices
  75   where range = (0, M.size graph - 1)
  76         pureGen = mkStdGen 137
  77         indices = take n $ unfoldr (Just . uniformR range) pureGen
  78
  79 getEdges :: [String] -> [(String, String)]
  80 getEdges xs = zipWith go xs (tail xs)
  81   where go a b
  82           | a < b = (a, b)
  83           | otherwise = (b, a)
  84
  85 removeEdge :: Graph -> (String, String) -> Graph
  86 removeEdge graph (a, b) = M.adjust (delete a) b $ M.adjust (delete b) a graph
  87
  88
  89 -- Parse the input file
  90
  91 graphP :: Parser Graph
  92 nodeP :: Parser (String, [String])
  93 nameP :: Parser String
  94
  95 graphP = M.fromList <$> nodeP `sepBy` endOfLine
  96 nodeP = (,) <$> nameP <* ": " <*> nameP `sepBy` " "
  97
  98 nameP = many1 letter
  99
 100 successfulParse :: Text -> Graph
 101 successfulParse input =
 102   case parseOnly graphP input of
 103     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 104     Right matches -> matches