Day7 Part 2

day7/main.hs

@@ -6,6 +6,7 @@
 {-# LANGUAGE OverloadedStrings #-}
 
 import Control.Applicative
+import Control.Monad.Zip
 import Data.Attoparsec.Text (Parser, parseOnly)
 import Data.Bool (bool)
 import Data.Map (Map)
@@ -13,6 +14,7 @@ import Data.Monoid (Sum(..))
 import Data.MultiSet (MultiSet)
 import Data.Tree (Tree)
 import Text.Pretty.Simple
+-- import Debug.Trace (trace)
 import qualified Data.Attoparsec.Text as A
 import qualified Data.Map.Strict as M
 import qualified Data.MultiSet as B
@@ -32,6 +34,18 @@ exampleData =
   , "dotted black bags contain no other bags."
   ]
 
+exampleData2 :: [String]
+exampleData2 =
+  [ "shiny gold bags contain 2 dark red bags."
+  , "dark red bags contain 2 dark orange bags."
+  , "dark orange bags contain 2 dark yellow bags."
+  , "dark yellow bags contain 2 dark green bags."
+  , "dark green bags contain 2 dark blue bags."
+  , "dark blue bags contain 2 dark violet bags."
+  , "dark violet bags contain no other bags."
+  ]
+
+
 data Color = Color String String
   deriving (Show, Ord, Eq)
 
@@ -107,6 +121,16 @@ solvePart1 color strs = do
   pure $ getSum $ foldMap (Sum . bool 0 1) containColor
   -- ^^ complicated exp for: pure $ length $ filter (== True) containColor
 
+updateTree :: Tree Int -> Tree Int
+updateTree t@(T.Node _ []) = t
+updateTree (T.Node v vx) = T.Node v (fmap ((fmap ((*) v)) . updateTree) vx)
+
+solvePart2 :: Color -> [String] -> Maybe ( Sum Int )
+solvePart2 color strs = do
+  rules <- parseStrs strs
+  let tree = mkTree (color,1) rules
+  pure $ foldMap Sum $ (updateTree . snd . munzip) tree
+
 main :: IO ()
 main = do
   input <- lines <$> readFile "day7/input"
@@ -124,3 +148,7 @@ main = do
   putStrLn ":: Solving Part 1"
   print $ solvePart1 (Color "shiny" "gold") exampleData
   print $ solvePart1 (Color "shiny" "gold") input
+  putStrLn ":: Solving Part 2"
+  print $ solvePart2 (Color "shiny" "gold") exampleData
+  print $ solvePart2 (Color "shiny" "gold") exampleData2
+  print $ solvePart2 (Color "shiny" "gold") input