adventofcode-2020/day7/main.hs

124 lines
3.8 KiB
Haskell
Raw Normal View History

2020-12-07 19:35:03 +01:00
#! /usr/bin/env -S"ANSWER=42" nix-shell
#! nix-shell -p ghcid
2020-12-08 16:17:47 +01:00
#! nix-shell -p "haskellPackages.ghcWithPackages (p: with p; [attoparsec multiset containers pretty-simple])"
2020-12-07 19:35:03 +01:00
#! nix-shell -i "ghcid -c 'ghci -Wall' -T main"
{-# LANGUAGE OverloadedStrings #-}
2020-12-08 16:17:47 +01:00
import Control.Applicative
import Data.Attoparsec.Text (Parser, parseOnly)
import Data.Map (Map)
import Data.MultiSet (MultiSet)
import Data.Tree (Tree)
import Text.Pretty.Simple
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
2020-12-07 21:07:42 +01:00
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."
]
2020-12-08 16:17:47 +01:00
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 $ length $ filter (== True) containColor
2020-12-07 19:35:03 +01:00
main :: IO ()
main = do
2020-12-08 16:17:47 +01:00
input <- lines <$> readFile "day7/input"
putStrLn ":: Day 7 - Part 1"
2020-12-07 21:07:42 +01:00
print $ exampleData
2020-12-08 16:17:47 +01:00
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