#!/usr/bin/env nix-shell
#!nix-shell --pure -i "runghc -- -i../" -p "haskellPackages.ghcWithPackages (pkgs: with pkgs; [ ])"
import Data.List
import qualified Data.Map as Map
import Text.Parsec
import Text.Parsec.Char
import Text.Parsec.Combinator
import Text.Parsec.String (Parser)
import Aoc (readAndParseStdin)
main = do
document <- readAndParseStdin parseDocument
print $ part1 document
print $ part2 document
part1 :: Document -> Int
part1 doc = traverseTreeUntilD doc (== "ZZZ") "AAA"
part2 :: Document -> Int
part2 doc = foldl1 lcm . map (traverseTreeUntilD doc ("Z" `isSuffixOf`)) $ startingPositions
where
startingPositions = filter ("A" `isSuffixOf`) $ (Map.keys . docMap) doc
traverseTreeUntilD :: Document -> (String -> Bool) -> String -> Int
traverseTreeUntilD doc = traverseTreeUntil (cycle $ directions doc) (docMap doc) 0
traverseTreeUntil :: [Direction] -> Map.Map String (String, String) -> Int -> (String -> Bool) -> String -> Int
traverseTreeUntil (x : xs) m n predicate elem
| predicate elem = n
| otherwise = traverseTreeUntil xs m (n + 1) predicate (readNext getNode x)
where
getNode = case Map.lookup elem m of
Just a -> a
Nothing -> error (show elem)
readNext (n, _) DirectionLeft = n
readNext (_, n) DirectionRight = n
parseDocument :: Parser Document
parseDocument = do
directions <- many1 parseDirection
_ <- string "\n\n"
nodes <- parseNode `endBy` char '\n'
return
Document
{ directions,
docMap = Map.fromList nodes
}
parseNode :: Parser (String, (String, String))
parseNode = do
node <- many1 alphaNum
_ <- spaces <* char '=' <* spaces
_ <- char '('
left <- many1 alphaNum
_ <- char ',' <* spaces
right <- many1 alphaNum
_ <- char ')'
return (node, (left, right))
parseDirection :: Parser Direction
parseDirection = do
c <- oneOf "LR"
return $ case c of
'L' -> DirectionLeft
'R' -> DirectionRight
data Document = Document
{ directions :: [Direction],
docMap :: Map.Map String (String, String)
}
deriving (Show)
data Direction = DirectionLeft | DirectionRight deriving (Enum, Show)