src/advent12/advent12.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2
   3 import Data.Text (Text)
   4 import qualified Data.Text as T
   5 import qualified Data.Text.IO as TIO
   6
   7 import Text.Megaparsec
   8 import qualified Text.Megaparsec.Lexer as L
   9 import Text.Megaparsec.Text (Parser)
  10
  11 import qualified Data.Map.Strict as M
  12 import Data.Map.Strict ((!))
  13
  14 import qualified Data.Set as S
  15 import Control.Applicative (empty)
  16
  17
  18 type ProgSet = S.Set Integer
  19 type Pipes = M.Map Integer ProgSet
  20
  21
  22 main :: IO ()
  23 main = do
  24         input <- TIO.readFile "data/advent12.txt"
  25         let pipes = successfulParse input
  26         print $ part1 pipes
  27         print $ part2 pipes
  28
  29
  30 part1 pipes = S.size $ reachable pipes (S.empty) (pipes!0)
  31
  32 part2 pipes = n
  33     where (_, n, _) = foldl addGroup (S.empty, 0, pipes) $ M.keys pipes
  34
  35
  36 addGroup :: (ProgSet, Integer, Pipes) -> Integer -> (ProgSet, Integer, Pipes)
  37 addGroup (done, n, pipes) p
  38     | p `S.member` done = (done, n, pipes)
  39     | otherwise = (S.union done reached, n + 1, pipes)
  40         where reached = reachable pipes (S.empty) (pipes!p)
  41
  42
  43 reachable :: Pipes -> ProgSet -> ProgSet -> ProgSet
  44 reachable pipes reached frontier
  45     | S.null frontier = reached
  46     | otherwise = reachable pipes reached' frontier'
  47         where frontier' = S.difference (unions' $ S.map (\p -> pipes!p) frontier) reached
  48               reached' = reached `S.union` frontier'
  49               unions' = S.foldl S.union S.empty
  50
  51
  52
  53 sc :: Parser ()
  54 sc = L.space (skipSome spaceChar) empty empty -- lineCmnt blockCmnt
  55   -- where
  56   --   lineCmnt  = L.skipLineComment "//"
  57   --   blockCmnt = L.skipBlockComment "/*" "*/"
  58
  59 lexeme  = L.lexeme sc
  60 integer = lexeme L.integer
  61 symb = L.symbol sc
  62
  63
  64 pipesP = many pipe
  65
  66 pipe = assocify <$> integer <*> (symb "<->" *> (integer `sepBy1` (symb ",")))
  67     where assocify a b = (a, S.fromList b)
  68
  69 successfulParse :: Text -> Pipes
  70 successfulParse input =
  71         case parse pipesP "input" input of
  72                 Left  err   -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
  73                 Right betterInput -> M.fromList betterInput
  74
  75
  76 -- sample = T.pack "0 <-> 2\n\
  77 -- \1 <-> 1\n\
  78 -- \2 <-> 0, 3, 4\n\
  79 -- \3 <-> 2, 4\n\
  80 -- \4 <-> 2, 3, 6\n\
  81 -- \5 <-> 6\n\
  82 -- \6 <-> 4, 5"