154 lines
4.9 KiB
Haskell
Executable file
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
|