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 import Data.Maybe
  16
  17 data QuantifiedBag = QuantifiedBag Integer String
  18     deriving (Show, Eq, Ord)
  19
  20
  21 main :: IO ()
  22 main =
  23   do  text <- TIO.readFile "data/advent07.txt"
  24       let bags = successfulParse text
  25       -- print bags
  26       -- print $ invertBags bags
  27       print $ part1 bags
  28       print $ part2 bags
  29
  30 part1 bags = S.size $ S.delete "shiny gold" containers
  31     where containers = bagsContaining (invertBags bags) (S.singleton "shiny gold") S.empty
  32
  33 part2 bags = (nContainedBags bags "shiny gold") - 1
  34
  35 invertBags bags = foldr addInvert M.empty $ concatMap swapPair $ M.assocs bags
  36     where swapPair (a, bs) = [(qName b, a) | b <- S.toList bs]
  37
  38 addInvert :: (String, String) -> (M.Map String (S.Set String)) -> (M.Map String (S.Set String))
  39 addInvert (k, v) m =
  40     if k `M.member` m
  41     then M.insert k v' m
  42     else M.insert k (S.singleton v) m
  43     where v' = S.insert v (m!k)
  44
  45 qName (QuantifiedBag _ n) = n
  46 qCount (QuantifiedBag n _) = n
  47
  48 bagsContaining :: (M.Map String (S.Set String)) -> (S.Set String) -> (S.Set String) -> (S.Set String)
  49 bagsContaining iBags agenda result
  50     | S.null agenda = result
  51     | otherwise = bagsContaining iBags agenda'' (S.insert thisColour result)
  52     where thisColour = S.findMin agenda
  53           agenda' = S.delete thisColour agenda
  54           agenda'' = if thisColour `S.member` result
  55                      then agenda'
  56                      else S.union (fromMaybe S.empty (M.lookup thisColour iBags)) agenda'
  57
  58 nContainedBags :: (M.Map String (S.Set QuantifiedBag)) -> String -> Integer
  59 nContainedBags bags thisBag = 1 + (sum $ map subCount others)
  60     where others = S.toList $ bags!thisBag
  61           subCount b = (qCount b) * (nContainedBags bags (qName b))
  62
  63
  64
  65 -- -- Parse the input file
  66
  67 bagNameP = manyTill anyChar ((string " bags") <|> (string " bag"))
  68
  69 quantifiedBagP = QuantifiedBag <$> decimal <* space <*> bagNameP
  70
  71 emptyBagListP = "no other bags" *> pure S.empty
  72
  73 bagListP = S.fromList <$> sepBy quantifiedBagP (string ", ")
  74
  75 bagContentsP = emptyBagListP <|> bagListP
  76
  77 ruleP = (,) <$> bagNameP <* " contain " <*> bagContentsP <* "."
  78
  79 rulesP = M.fromList <$> sepBy ruleP endOfLine
  80
  81
  82
  83 -- successfulParse :: Text -> [[S.Set Char]]
  84 successfulParse input =
  85   case parseOnly rulesP input of
  86     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
  87     Right bags -> bags