Day07/Main.hs

   1 {-# LANGUAGE OverloadedStrings #-} -- for building Text strings
   2 module Main where
   3 import Control.Applicative (Alternative((<|>)))
   4 import Control.Monad (void, when)
   5 import Data.Text.Lazy (Text)
   6 import Data.Void (Void)
   7 import Prelude
   8 import qualified Data.Text.Lazy.Encoding as Text
   9 import qualified Text.Megaparsec as P
  10 import qualified Text.Megaparsec.Char as P
  11 import qualified Text.Megaparsec.Char.Lexer as P
  12 -- import qualified Text.Megaparsec.Debug as P
  13 import qualified Data.ByteString.Lazy as BSL
  14 import Paths_AoC2020
  15 import Data.Map.Strict (Map)
  16 import qualified Data.Map.Strict as Map
  17
  18 -- | See https://adventofcode.com/2020/day/7 for the problem statements
  19 data Day07Results = Day07Results
  20   { example1 :: Int
  21   , batch1 :: Int
  22   , example2 :: Int
  23   , example3 :: Int
  24   , batch2 :: Int
  25   } deriving (Show)
  26
  27 main :: IO ()
  28 main = do
  29   putStr "Day07Inputs: " >> getDataFileName "" >>= putStrLn
  30   print =<< Day07Results
  31     <$> (countTransitiveParentBags "shiny gold" <$> parse "example1")
  32     <*> (countTransitiveParentBags "shiny gold" <$> parse "batch")
  33     <*> (countChildrenBags "shiny gold" <$> parse "example1")
  34     <*> (countChildrenBags "shiny gold" <$> parse "example2")
  35     <*> (countChildrenBags "shiny gold" <$> parse "batch")
  36
  37 type Parser output = P.Parsec {-error-}Void {-input-}Text output
  38
  39 type Bag = Text
  40 type Rule = (Bag, Map Bag Int)
  41 type Rules = Map Bag (Map Bag Int)
  42
  43 parserBag :: Parser Bag
  44 parserBag =
  45   (fst <$>) $ P.match $ do
  46     void $ P.takeWhileP (Just "adj") (/=' ')
  47     void $ P.char ' '
  48     void $ P.takeWhileP (Just "col") (/=' ')
  49
  50 parserRule :: Parser Rule
  51 parserRule = {- P.dbg "parserRule" $ -} do
  52   outBag <- parserBag
  53   void $ P.string " bags contain "
  54   inBags <- P.choice
  55     [ [] <$ P.string "no other bags"
  56     , (`P.sepBy` P.string ", ") $ do
  57         cnt <- P.decimal
  58         void $ P.char ' '
  59         bag <- parserBag
  60         void $ P.string " bag"
  61         when (cnt > 1) $
  62           void $ P.optional (P.char 's')
  63         return (bag, cnt)
  64     ]
  65   void $ P.char '.'
  66   return (outBag, Map.fromList inBags)
  67
  68 parserRules :: Parser Rules
  69 parserRules =
  70   -- P.dbg "parserRules" $
  71   (Map.fromList <$>) $
  72   P.many $ parserRule <* (() <$ P.char '\n' <|> P.eof)
  73
  74 parse :: FilePath -> IO Rules
  75 parse input = do
  76   content <- Text.decodeUtf8 <$> (BSL.readFile =<< getDataFileName input)
  77   case P.parse (parserRules <* P.eof) input content of
  78     Left err -> error (P.errorBundlePretty err)
  79     Right rules -> return rules
  80
  81 type ReverseRules = Map Bag (Map Bag ())
  82
  83 reverseRules :: Rules -> ReverseRules
  84 reverseRules =
  85   Map.foldlWithKey (\acc outBag inBags ->
  86     Map.unionWith (<>) acc (Map.singleton outBag () <$ inBags)
  87   ) Map.empty
  88
  89 parentBags :: Bag -> ReverseRules -> Map Bag ()
  90 parentBags = Map.findWithDefault Map.empty
  91
  92 transitiveParentBags :: ReverseRules -> Bag -> Map Bag ()
  93 transitiveParentBags revRules bag =
  94   parentBags bag revRules <>
  95   Map.unions (transitiveParentBags revRules <$> Map.keys (parentBags bag revRules))
  96
  97 countTransitiveParentBags :: Bag -> Rules -> Int
  98 countTransitiveParentBags bag rules = Map.size $ transitiveParentBags (reverseRules rules) bag
  99
 100 countChildrenBags :: Bag -> Rules -> Int
 101 countChildrenBags bag rules =
 102   sum $ Map.mapWithKey (\childBag childBagCount ->
 103     childBagCount + childBagCount * countChildrenBags childBag rules
 104   ) childrenBags
 105   where childrenBags = Map.findWithDefault Map.empty bag rules