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 qualified Data.Map.Strict as M
  14 import Data.Map.Strict ((!), (\\))
  15
  16 -- from satellite to primary
  17 type Orbits = M.Map String String
  18
  19 -- transfer steps to each primary
  20 type TransferDistances = M.Map String Int
  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 directOrbits
  29         print $ part2 orbits
  30
  31 part1 :: Orbits -> [(String, String)] -> Int
  32 part1 orbits directOrbits = sum $ map (orbitCount orbits) satellites
  33     where satellites = map snd directOrbits
  34
  35 part2 :: Orbits -> Int
  36 part2 orbits = youDist + sanDist
  37     where youTrans = transferDistance orbits M.empty (orbits!"YOU") 0
  38           sanTrans = transferDistance orbits M.empty (orbits!"SAN") 0
  39           onlyYou = youTrans \\ sanTrans
  40           onlySan = sanTrans \\ youTrans
  41           -- youDist = 1 + (maximum $ M.elems onlyYou)
  42           -- sanDist = 1 + (maximum $ M.elems onlySan)
  43           youDist = M.size onlyYou
  44           sanDist = M.size onlySan
  45
  46
  47 buildOrbits :: [(String, String)] -> Orbits
  48 buildOrbits = foldl' addOrbit M.empty
  49
  50 addOrbit :: Orbits -> (String, String) -> Orbits
  51 addOrbit orbits (primary, satellite) = M.insert satellite primary orbits
  52
  53 orbitCount :: Orbits -> String -> Int
  54 orbitCount orbits here
  55     | here `M.member` orbits = 1 + (orbitCount orbits (orbits!here))
  56     | otherwise = 0
  57
  58 transferDistance :: Orbits -> TransferDistances -> String -> Int -> TransferDistances
  59 transferDistance orbits transfers here dist
  60     | here `M.member` orbits = transferDistance orbits transfers' there (dist + 1)
  61     | otherwise = transfers'
  62     where there = orbits!here
  63           transfers' = M.insert here dist transfers
  64
  65
  66 -- Parse the input file
  67 type Parser = Parsec Void Text
  68
  69 sc :: Parser ()
  70 sc = L.space (skipSome spaceChar) CA.empty CA.empty
  71 -- sc = L.space (skipSome (char ' ')) CA.empty CA.empty
  72
  73 -- lexeme  = L.lexeme sc
  74 -- integer = lexeme L.decimal
  75 -- signedInteger = L.signed sc integer
  76 symb = L.symbol sc
  77 orbSep = symb ")"
  78 identifierP = some alphaNumChar <* sc
  79
  80 orbitsP = many orbitP
  81 orbitP = (,) <$> identifierP <* orbSep <*> identifierP
  82
  83 successfulParse :: Text -> [(String, String)]
  84 successfulParse input =
  85         case parse orbitsP "input" input of
  86                 Left  _err -> [] -- TIO.putStr $ T.pack $ parseErrorPretty err
  87                 Right orbits -> orbits