advent08/Main.hs

   1 -- Writeup at https://work.njae.me.uk/2023/12/08/advent-of-code-2023-day-8/
   2
   3 import AoC
   4 import Data.Text (Text)
   5 import qualified Data.Text.IO as TIO
   6 import Data.Attoparsec.Text -- hiding (take)
   7 import Control.Applicative
   8 import Data.List
   9 import qualified Data.Map.Strict as M
  10 import Data.Map.Strict ((!), (!?))
  11
  12 data Direction = L | R deriving (Show, Eq)
  13 data Node = Node String String deriving (Show)
  14 type Desert = M.Map String Node
  15
  16 data State = State { getHere :: String, getSteps :: Int } deriving (Eq, Show)
  17 instance Ord State where
  18   compare (State n1 s1) (State n2 s2) = compare (s1, n1) (s2, n2)
  19
  20 data PathStart = PathStart String Int deriving (Eq, Ord, Show)
  21 data CollapsedPath = CPath String Int deriving (Show)
  22 type Paths = M.Map PathStart CollapsedPath
  23
  24 main :: IO ()
  25 main =
  26   do  dataFileName <- getDataFileName
  27       text <- TIO.readFile dataFileName
  28       let (directions, desert) = successfulParse text
  29       -- print $ length directions
  30       -- print $ directions
  31       -- print $ desert
  32       print $ part1 desert directions
  33       print $ part2 desert directions
  34       print $ part3 desert directions
  35
  36 part1, part2 :: Desert -> [Direction] -> Int
  37 part1 desert directions = getSteps $ walk desert directions (State "AAA" 0)
  38
  39
  40 part2 desert directions = foldl1 lcm pathLens
  41   where cache = generateRouteLengths desert directions
  42         pathLens = fmap (\(CPath _ l) -> l) $ M.elems cache
  43         -- directionsLen = length directions
  44
  45 generateRouteLengths :: Desert -> [Direction] -> Paths
  46 generateRouteLengths desert directions = M.unions ((fmap snd sResults) ++ (fmap snd gResults))
  47   where starts = sort $ fmap (\s -> State s 0) $ startsOf desert
  48         sResults = fmap (walkWithCache desert directions M.empty) starts
  49         fromGoals = fmap fst sResults
  50         gResults = fmap (walkWithCache desert directions M.empty) fromGoals
  51
  52
  53 part3 desert directions = multiWalk desert directions M.empty starts
  54   where starts = fmap (\s -> State s 0) $ startsOf desert
  55
  56 multiWalk desert directions cache states@(s:ss)
  57   | (all isGoal states) && (sameTime states) = states
  58   | otherwise = multiWalk desert directions newCache $ sort (s':ss)
  59   where (s', newCache) = walkWithCache desert directions cache s
  60
  61 sameTime states = (length $ nub times) == 1
  62   where times = fmap getSteps states
  63
  64 walk :: Desert -> [Direction] -> State -> State
  65 walk desert directions start = head $ dropWhile (not . isGoal) path
  66   where path = scanl' (step desert) start $ drop offset $ cycle directions
  67         offset = (getSteps start) `mod` (length directions)
  68
  69 step :: Desert -> State -> Direction -> State
  70 step desert (State here steps) direction
  71   | direction == L = State thereL (steps + 1)
  72   | direction == R = State thereR (steps + 1)
  73   where
  74     (Node thereL thereR) = desert ! here
  75
  76
  77 startsOf :: Desert -> [String]
  78 startsOf desert = filter ((=='A') . last) $ M.keys desert
  79 -- goalsOf  desert = filter ((=='Z') . last) $ M.keys desert
  80
  81 isGoal :: State -> Bool
  82 isGoal (State here _) = (last here) == 'Z'
  83
  84 walkWithCache :: Desert -> [Direction] -> Paths -> State -> (State, Paths)
  85 walkWithCache desert directions cache start =
  86   case cacheEntry of
  87     Just (CPath there steps) -> (State there (steps + (getSteps start)), cache)
  88     Nothing -> (newState, newCache)
  89   where offset = (getSteps start) `mod` (length directions)
  90         cacheEntry = cache !? (PathStart (getHere start) offset)
  91         start' = step desert start (directions !! offset)
  92         newState = walk desert directions start'
  93         newCache = M.insert (PathStart (getHere start) offset)
  94                             (CPath (getHere newState) ((getSteps newState) - (getSteps start)))
  95                             cache
  96
  97 -- Parse the input file
  98
  99 problemP :: Parser ([Direction], Desert)
 100 directionP :: Parser Direction
 101 desertP :: Parser Desert
 102 desertLineP :: Parser (String, Node)
 103 nodeP :: Parser Node
 104 nameP :: Parser String
 105
 106 problemP = (,) <$> ((many1 directionP) <* many1 endOfLine) <*> desertP
 107 directionP = (L <$ "L") <|> (R <$ "R")
 108
 109 desertP = M.fromList <$> desertLineP `sepBy` endOfLine
 110 desertLineP = (,) <$> (nameP <* " = ") <*> nodeP
 111
 112 nodeP = Node <$> ("(" *> nameP <* ", ") <*> (nameP <* ")")
 113 nameP = many1 (letter <|> digit)
 114
 115 successfulParse :: Text -> ([Direction], Desert)
 116 successfulParse input =
 117   case parseOnly problemP input of
 118     Left  _err -> ([], M.empty) -- TIO.putStr $ T.pack $ parseErrorPretty err
 119     Right matches -> matches