advent07/src/advent07graph.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.Graph.DGraph as D
  12 import qualified Data.Graph.Types as G
  13 import qualified Data.Graph.Traversal as T
  14
  15 import qualified Data.Set as S
  16 import qualified Data.Map.Strict as M
  17
  18 data QuantifiedBag = QuantifiedBag Int String
  19     deriving (Show, Eq, Ord)
  20
  21 type QBags = S.Set QuantifiedBag
  22 type BagRules = M.Map String QBags
  23 type BagGraph = D.DGraph String Int
  24
  25
  26 main :: IO ()
  27 main =
  28   do  text <- TIO.readFile "data/advent07.txt"
  29       let bags = successfulParse text
  30       let graph = buildGraph bags
  31       -- print graph
  32       -- dumpBagDot bags
  33       print $ part1 graph
  34       print $ part2 graph
  35
  36 part1 graph = length (T.bfsVertices (D.transpose graph) "shiny gold") - 1
  37
  38 part2 graph = (bfsCount graph "shiny gold") - 1
  39
  40 bfsCount :: BagGraph -> String -> Int
  41 bfsCount graph thisBag = 1 + (sum $ map subCount others)
  42     where others = D.outboundingArcs graph thisBag
  43           subCount a = (G.attribute a) * (bfsCount graph $ G.destinationVertex a)
  44
  45
  46 buildGraph :: BagRules -> BagGraph
  47 buildGraph rules = M.foldrWithKey addRule G.empty rules
  48
  49 addRule :: String -> QBags -> BagGraph -> BagGraph
  50 addRule source dests graph = S.foldr (addArc source) graph dests
  51
  52 addArc :: String -> QuantifiedBag -> BagGraph -> BagGraph
  53 addArc source (QuantifiedBag quantity destination) graph = D.insertArc arc graph
  54   where arc = G.Arc source destination quantity
  55
  56
  57 -- dumpBagDot bags =
  58 --     do writeFile "a07dump.dot" "digraph {\n"
  59 --        mapM_ dumpABag (M.assocs bags)
  60 --        appendFile "a07dump.dot" "shiny_gold [fillcolor = gold1 ]\n"
  61 --        appendFile "a07dump.dot" "}\n"
  62
  63 -- dumpABag (bag, contents) =
  64 --     mapM_ (dumpALink bag) (S.toList contents)
  65
  66 -- dumpALink bag (QuantifiedBag n name) =
  67 --     do  let name' = squashName name
  68 --         let bag' = squashName bag
  69 --         let txt = bag' ++ " -> " ++ name' ++ "\n"
  70 --         appendFile "a07dump.dot" txt
  71
  72 -- squashName :: String -> String
  73 -- squashName name = [if c == ' ' then '_' else c | c <- name]
  74
  75
  76 -- -- Parse the input file
  77
  78 bagNameP = manyTill anyChar ((string " bags") <|> (string " bag"))
  79
  80 quantifiedBagP = QuantifiedBag <$> decimal <* space <*> bagNameP
  81
  82 emptyBagListP = "no other bags" *> pure S.empty
  83
  84 bagListP = S.fromList <$> sepBy quantifiedBagP (string ", ")
  85
  86 bagContentsP = emptyBagListP <|> bagListP
  87
  88 ruleP = (,) <$> bagNameP <* " contain " <*> bagContentsP <* "."
  89
  90 rulesP = M.fromList <$> sepBy ruleP endOfLine
  91
  92
  93 successfulParse :: Text -> BagRules
  94 successfulParse input =
  95   case parseOnly rulesP input of
  96     Left  _err -> M.empty -- TIO.putStr $ T.pack $ parseErrorPretty err
  97     Right bags -> bags