Our first goal here is to simplify the input. We conjure up the replace function and make some replacements to our
text. We also define a data type for bags holding other bags and how to parse the cleaned up text.
data Relation = Relation String (HashSet (Int, String))
deriving (Show, Eq)
replace :: Eq a => [a] -> [a] -> [a] -> [a]
replace n r h = intercalate r (splitOn n h)
makeRelation :: [String] -> Relation
makeRelation (x:xs) = Relation x (relations xs)
where relations (a:b:rest) = singleton (read a, b) <> relations rest
relations [] = mempty
relations _ = error "impossibru"
makeRelation _ = error "impossible"
relations = fmap makeRelation . cleanup' . cleanup $ input
cleanup = replace "," "" . replace "." "" . replace "bag" "" . replace "bags" "bag" . replace "contain" ""
cleanup' = fmap (fmap concat . split (whenElt (all isDigit)) . words) . lines
Here whenElt is a exceedingly poorly named function (in my opinion) which will split when the condition is met.
Therefore cleanup' will split the input into lines, and then split each line into words and group them into [color,
count, color, count, color, ...].
Part A wants us to traverse the relation backwards, so we just make a forwards tree and BFS forwards instead.
type UpGraph = HashMap String (HashSet String)
(<>?) :: Monoid a => ASetter' s (Maybe a) -> a -> s -> s
l <>? a = over l (fmap (<>a) . maybe (Just mempty) Just)
buildContainer :: UpGraph -> Relation -> UpGraph
buildContainer g (Relation up rs) = foldl' (\g' r -> g' & at (snd r) <>? singleton up) g rs
bfsA :: String -> HashSet String -> UpGraph -> HashSet String
bfsA start visited g =
if start `member` visited then mempty
else
case g ^? ix start of
Just ups -> ups <> foldl' (<>) mempty (HS.map (\u -> bfsA u visited' g) ups)
Nothing -> singleton start
where visited' = visited <> singleton start
container = foldl' buildContainer mempty $ relations
Here I define this nifty little combinator combining both ?~ and <>~ which will let me use at and transform
Nothing into mempty. This lets us avoid the case where the key does not yet exist. The BFS is fairly straightfoward.
It is important to remember that you also need to count the bags along the way instead of only the bags at the end. The
answer is then the length of the resulting set length $ bfsA "shinygold" mempty container.
Part B is a bit more straightforward. We again define a different type of tree and perform a simple BFS.
type DownGraph = HashMap String (HashSet (Int, String))
buildContained :: DownGraph -> Relation -> DownGraph
buildContained g (Relation up rs) = g & at up ?~ rs
bfsB :: String -> HashSet String -> DownGraph -> Int
bfsB start visited g =
if start `member` visited then 0
else
case g ^? ix start of
Just rs -> 1 + sum ((\(n, r) -> n * bfsB r visited' g) <$> toList rs)
Nothing -> 1
where visited' = visited <> singleton start
contained = foldl' buildContained mempty $ relations
Here the BFS returns the number of sub bags including itself. Therefore we must be careful and subtract one to get our
final answer of bfsB "shinygold" mempty contained - 1. We can, however, also write the BFS as
case g ^? ix start of
Just rs -> sum ((\(n, r) -> n + n * bfsB r visited' g) <$> toList rs)
Nothing -> 0
which will directly the sub bags at the cost of a few more additions if subtracting one is not groovy enough for you.