advent07/src/advent07.hs

   1 -- import Debug.Trace
   2
   3 import Data.Text (Text)
   4 import qualified Data.Text as T
   5 import qualified Data.Text.IO as TIO
   6
   7 import Data.Attoparsec.Text
   8 -- import Data.Attoparsec.Combinator
   9 import Control.Applicative
  10
  11 import qualified Data.Set as S
  12 import qualified Data.Map.Strict as M
  13 import Data.Map.Strict ((!))
  14
  15 data QuantifiedBag = QuantifiedBag Integer String
  16     deriving (Show, Eq, Ord)
  17
  18 qName (QuantifiedBag _ n) = n
  19 qCount (QuantifiedBag n _) = n
  20
  21 type Bags = S.Set String
  22 type QBags = S.Set QuantifiedBag
  23 type BagRules = M.Map String QBags
  24 type InvertBags = M.Map String Bags
  25
  26
  27 main :: IO ()
  28 main =
  29   do  text <- TIO.readFile "data/advent07.txt"
  30       let bags = successfulParse text
  31       -- dumpBagDot bags
  32       print $ part1 bags
  33       print $ part2 bags
  34
  35 -- dumpBagDot bags =
  36 --     do writeFile "a07dump.dot" "digraph {\n"
  37 --        mapM_ dumpABag (M.assocs bags)
  38 --        appendFile "a07dump.dot" "shiny_gold [fillcolor = gold1 ]\n"
  39 --        appendFile "a07dump.dot" "}\n"
  40
  41 -- dumpABag (bag, contents) =
  42 --     mapM_ (dumpALink bag) (S.toList contents)
  43
  44 -- dumpALink bag (QuantifiedBag n name) =
  45 --     do  let name' = squashName name
  46 --         let bag' = squashName bag
  47 --         let txt = bag' ++ " -> " ++ name' ++ "\n"
  48 --         appendFile "a07dump.dot" txt
  49
  50 -- squashName :: String -> String
  51 -- squashName name = [if c == ' ' then '_' else c | c <- name]
  52
  53
  54 part1 bags = S.size $ S.delete "shiny gold" containers
  55     where containers = bagsContaining (invertBags bags) (S.singleton "shiny gold") S.empty
  56
  57 part2 bags = (nContainedBags bags "shiny gold") - 1
  58
  59 invertBags :: BagRules -> InvertBags
  60 invertBags bags = foldr addInvert M.empty $ concatMap swapPair $ M.assocs bags
  61     where swapPair (a, bs) = [(qName b, a) | b <- S.toList bs]
  62
  63 addInvert :: (String, String) -> InvertBags -> InvertBags
  64 addInvert (k, v) m = M.insert k v' m
  65     where v' = S.insert v (M.findWithDefault S.empty k m)
  66
  67
  68 bagsContaining :: InvertBags -> Bags -> Bags -> Bags
  69 bagsContaining iBags agenda result
  70     | S.null agenda = result
  71     | otherwise = bagsContaining iBags agenda'' (S.insert thisColour result)
  72     where thisColour = S.findMin agenda
  73           agenda' = S.delete thisColour agenda
  74           agenda'' = if thisColour `S.member` result
  75                      then agenda'
  76                      else S.union (M.findWithDefault S.empty thisColour iBags) agenda'
  77
  78 nContainedBags :: BagRules -> String -> Integer
  79 nContainedBags bags thisBag = 1 + (sum $ map subCount others)
  80     where others = S.toList $ bags!thisBag
  81           subCount b = (qCount b) * (nContainedBags bags (qName b))
  82
  83
  84
  85 -- -- Parse the input file
  86
  87 bagNameP = manyTill anyChar ((string " bags") <|> (string " bag"))
  88
  89 quantifiedBagP = QuantifiedBag <$> decimal <* space <*> bagNameP
  90
  91 emptyBagListP = "no other bags" *> pure S.empty
  92
  93 bagListP = S.fromList <$> sepBy quantifiedBagP (string ", ")
  94
  95 bagContentsP = emptyBagListP <|> bagListP
  96
  97 ruleP = (,) <$> bagNameP <* " contain " <*> bagContentsP <* "."
  98
  99 rulesP = M.fromList <$> sepBy ruleP endOfLine
 100
 101
 102 successfulParse :: Text -> BagRules
 103 successfulParse input =
 104   case parseOnly rulesP input of
 105     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
 106     Right bags -> bags