1 import Data.Text (Text)
2 import qualified Data.Text.IO as TIO
4 import Data.Void (Void)
6 import Text.Megaparsec hiding (State)
7 import Text.Megaparsec.Char
8 import qualified Text.Megaparsec.Char.Lexer as L
9 import qualified Control.Applicative as CA
11 import Data.List (foldl')
12 import qualified Data.Set as S
13 import Data.Set ((\\))
14 import qualified Data.Map.Strict as M
15 import Data.Map.Strict ((!))
17 -- from satellite to primary
18 type Orbits = M.Map String String
20 type Transfers = S.Set String
25 text <- TIO.readFile "data/advent06.txt"
26 let directOrbits = successfulParse text
27 let orbits = buildOrbits directOrbits
31 part1 :: Orbits -> Int
32 part1 orbits = sum $ map (orbitCount orbits) $ M.keys orbits
34 part2 :: Orbits -> Int
35 part2 orbits = youDist + sanDist
36 where youTrans = transferSteps orbits S.empty (orbits!"YOU")
37 sanTrans = transferSteps orbits S.empty (orbits!"SAN")
38 onlyYou = youTrans \\ sanTrans
39 onlySan = sanTrans \\ youTrans
40 youDist = S.size onlyYou
41 sanDist = S.size onlySan
44 buildOrbits :: [(String, String)] -> Orbits
45 buildOrbits = foldl' addOrbit M.empty
47 addOrbit :: Orbits -> (String, String) -> Orbits
48 addOrbit orbits (primary, satellite) = M.insert satellite primary orbits
50 orbitCount :: Orbits -> String -> Int
51 orbitCount orbits here
52 | here `M.member` orbits = 1 + (orbitCount orbits (orbits!here))
55 transferSteps :: Orbits -> Transfers -> String -> Transfers
56 transferSteps orbits transfers here
57 | here `M.member` orbits = transferSteps orbits transfers' there
58 | otherwise = transfers'
59 where there = orbits!here
60 transfers' = S.insert here transfers
63 -- Parse the input file
64 type Parser = Parsec Void Text
67 sc = L.space (skipSome spaceChar) CA.empty CA.empty
68 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
70 -- lexeme = L.lexeme sc
71 -- integer = lexeme L.decimal
72 -- signedInteger = L.signed sc integer
75 identifierP = some alphaNumChar <* sc
78 orbitP = (,) <$> identifierP <* orbSep <*> identifierP
80 successfulParse :: Text -> [(String, String)]
81 successfulParse input =
82 case parse orbitsP "input" input of
83 Left _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
84 Right orbits -> orbits