advent06/src/advent06.hs

   1 import Data.Text (Text)
   2 import qualified Data.Text.IO as TIO
   3
   4 import Data.Void (Void)
   5
   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
  10
  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 ((!))
  16
  17 -- from satellite to primary
  18 type Orbits = M.Map String String
  19
  20 type Transfers = S.Set String
  21
  22
  23 main :: IO ()
  24 main = do
  25         text <- TIO.readFile "data/advent06.txt"
  26         let directOrbits = successfulParse text
  27         let orbits = buildOrbits directOrbits
  28         print $ part1 orbits
  29         print $ part2 orbits
  30
  31 part1 :: Orbits -> Int
  32 part1 orbits = sum $ map (orbitCount orbits) $ M.keys orbits
  33
  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
  42
  43
  44 buildOrbits :: [(String, String)] -> Orbits
  45 buildOrbits = foldl' addOrbit M.empty
  46
  47 addOrbit :: Orbits -> (String, String) -> Orbits
  48 addOrbit orbits (primary, satellite) = M.insert satellite primary orbits
  49
  50 orbitCount :: Orbits -> String -> Int
  51 orbitCount orbits here
  52     | here `M.member` orbits = 1 + (orbitCount orbits (orbits!here))
  53     | otherwise = 0
  54
  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
  61
  62
  63 -- Parse the input file
  64 type Parser = Parsec Void Text
  65
  66 sc :: Parser ()
  67 sc = L.space (skipSome spaceChar) CA.empty CA.empty
  68 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
  69
  70 -- lexeme  = L.lexeme sc
  71 -- integer = lexeme L.decimal
  72 -- signedInteger = L.signed sc integer
  73 symb = L.symbol sc
  74 orbSep = symb ")"
  75 identifierP = some alphaNumChar <* sc
  76
  77 orbitsP = many orbitP
  78 orbitP = (,) <$> identifierP <* orbSep <*> identifierP
  79
  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