adventofcode-2020/day7/main.hs
2020-12-08 22:21:45 +02:00

154 lines
4.9 KiB
Haskell
Executable file

#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [attoparsec multiset containers pretty-simple])"
#! nix-shell -i "ghcid -c 'ghci -Wall' -T main"
{-# LANGUAGE OverloadedStrings #-}
import Control.Applicative
import Control.Monad.Zip
import Data.Attoparsec.Text (Parser, parseOnly)
import Data.Bool (bool)
import Data.Map (Map)
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
import qualified Data.Text as Text
import qualified Data.Tree as T
exampleData :: [String]
exampleData =
[ "light red bags contain 1 bright white bag, 2 muted yellow bags."
, "dark orange bags contain 3 bright white bags, 4 muted yellow bags."
, "bright white bags contain 1 shiny gold bag."
, "muted yellow bags contain 2 shiny gold bags, 9 faded blue bags."
, "shiny gold bags contain 1 dark olive bag, 2 vibrant plum bags."
, "dark olive bags contain 3 faded blue bags, 4 dotted black bags."
, "vibrant plum bags contain 5 faded blue bags, 6 dotted black bags."
, "faded blue bags contain no other bags."
, "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)
-- An efficient implementation of multisets, also sometimes called bags.
type Bag = MultiSet Color
parseFullColor :: Parser Color
parseFullColor = do
adjective <- A.many' A.letter
A.skipSpace
colorName <- A.many' A.letter
A.skipSpace
_ <- ("bags" <|> "bag")
pure $ Color adjective colorName
parseBag' :: Parser Bag
parseBag' = do
i <- A.decimal
A.skipSpace
c <- parseFullColor
pure $ B.insertMany c i B.empty
parseNoBag :: Parser Bag
parseNoBag = do
_ <- "no other bags"
pure $ B.empty
parseBag :: Parser Bag
parseBag = parseNoBag <|> parseBag'
type Rules = Map Color Bag
parseRule :: Parser Rules
parseRule = do
k <- parseFullColor
A.skipSpace
_ <- "contain"
A.skipSpace
vx <- parseBag `A.sepBy` (", ")
_ <- "."
let v = B.unions vx
pure $ M.insert k v M.empty
eitherToMaybe :: Either b a -> Maybe a
eitherToMaybe (Right v) = Just v
eitherToMaybe (Left _) = Nothing
parseStrs :: [String] -> Maybe Rules
parseStrs s = M.unions <$> parse s
where
parse = eitherToMaybe
. sequence
. map (parseOnly (parseRule <* A.endOfInput) . Text.pack)
mkTree :: (Color,Int) -> Rules -> Tree (Color, Int)
mkTree color rules = T.unfoldTree mkNode color
where
mkNode :: (Color, Int) -> ((Color, Int), [(Color, Int)])
mkNode n@(c,_) = (n, concat $ B.toOccurList <$> M.lookup c rules)
mkSimpleTree :: Color -> Rules -> Tree Color
mkSimpleTree color rules = T.unfoldTree mkNode color
where
mkNode :: Color -> (Color, [Color])
mkNode c = (c, fmap fst <$> concat $ B.toOccurList <$> M.lookup c rules)
solvePart1 :: Color -> [String] -> Maybe Int
solvePart1 color strs = do
rules <- parseStrs strs
let trees = map (\c -> mkSimpleTree c rules) (filter (/= color) $ M.keys rules)
let containColor = map (elem color) trees
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"
putStrLn ":: Day 7 - Part 1"
print $ exampleData
putStrLn ":: Tests"
print $ parseOnly (parseFullColor <* A.endOfInput) "light red bag"
print $ parseOnly (parseBag <* A.endOfInput) "2 muted blue bags"
print $ parseOnly (parseBag <* A.endOfInput) "3 muted red bages"
print $ parseOnly (parseRule <* A.endOfInput) "dark orange bags contain 3 bright white bags, 4 muted yellow bags."
putStrLn ":: Parsing test data"
print $ parseStrs exampleData
putStrLn ":: Building tree"
pPrint $ mkTree (Color "shiny" "gold", 1) <$> (parseStrs exampleData)
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