From a97a5be6c8c0ea5050351f0c0820cd7ec0bffe89 Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Mon, 7 Dec 2020 10:00:14 +0000
Subject: [PATCH] [hs] Clean up 2020_07

---
 hs/package.yaml         |  1 +
 hs/src/Aoc/Y2020/D07.hs | 45 +++++++++++++++++++++--------------------
 2 files changed, 24 insertions(+), 22 deletions(-)

diff --git a/hs/package.yaml b/hs/package.yaml
index d6d4646..6f35e39 100644
--- a/hs/package.yaml
+++ b/hs/package.yaml
@@ -7,6 +7,7 @@ dependencies:
 - megaparsec
 - optparse-applicative
 - text
+- transformers
 
 library:
   source-dirs: src
diff --git a/hs/src/Aoc/Y2020/D07.hs b/hs/src/Aoc/Y2020/D07.hs
index a0e89a3..57bb9d7 100644
--- a/hs/src/Aoc/Y2020/D07.hs
+++ b/hs/src/Aoc/Y2020/D07.hs
@@ -4,12 +4,13 @@ module Aoc.Y2020.D07
   ( day
   ) where
 
-import Control.Monad
-import Data.Maybe
+import           Control.Monad
+import           Data.Maybe
 
-import qualified Data.Map  as Map
-import qualified Data.Set  as Set
-import qualified Data.Text as T
+import           Control.Monad.Trans.State
+import qualified Data.Map                  as Map
+import qualified Data.Set                  as Set
+import qualified Data.Text                 as T
 
 import           Aoc.Day
 import           Aoc.Parse
@@ -18,26 +19,19 @@ type BagName = T.Text
 
 data BagCount = BagCount
   { bcName :: BagName
-  , bcNum :: Int
+  , bcNum  :: Int
   } deriving (Show)
 
 type BagMap = Map.Map BagName [BagCount]
 
 pBagName :: Parser T.Text
 pBagName = do
-  first <- word
-  void $ char ' '
-  second <- word
-  void $ char ' '
-  void $ string "bags" <|> string "bag"
-  pure $ first <> " " <> second
+  name <- mconcat <$> sequenceA [word, string " ", word]
+  void $ string " bags" <|> string " bag"
+  pure name
 
 pBagCount :: Parser BagCount
-pBagCount = do
-  n <- decimal
-  void $ char ' '
-  name <- pBagName
-  pure $ BagCount name n
+pBagCount = flip BagCount <$> (decimal <* char ' ') <*> pBagName
 
 parser :: Parser BagMap
 parser = fmap Map.fromList $ manyLines $ do
@@ -47,10 +41,16 @@ parser = fmap Map.fromList $ manyLines $ do
   void $ char '.'
   pure (name, bags)
 
-children :: BagMap -> BagName -> Set.Set BagName
-children m b =
-  let ch = map bcName $ fromMaybe [] $ m Map.!? b
-  in  foldr Set.union (Set.fromList ch) $ map (children m) ch
+children :: BagMap -> BagName -> State (Map.Map BagName (Set.Set BagName)) (Set.Set BagName)
+children m b = do
+  sm <- get
+  case sm Map.!? b of
+    Just result -> pure result
+    Nothing -> do
+      let ch = map bcName $ fromMaybe [] $ m Map.!? b
+      result <- Set.unions . (Set.fromList ch :) <$> traverse (children m) ch
+      put $ Map.insert b result sm
+      pure result
 
 countChildren :: BagMap -> BagName -> Int
 countChildren m b =
@@ -63,7 +63,8 @@ myBag = "shiny gold"
 solver :: BagMap -> IO ()
 solver bags = do
   putStrLn ">> Part 1"
-  print $ length $ filter (Set.member myBag) $ map (children bags) $ Set.toList $ Map.keysSet bags
+  let childrenPerBag = flip evalState Map.empty $ traverse (children bags) $ Map.keys bags
+  print $ length $ filter (Set.member myBag) childrenPerBag
 
   putStrLn ""
   putStrLn ">> Part 2"