-
Notifications
You must be signed in to change notification settings - Fork 5
/
Day07.hs
62 lines (54 loc) · 2.02 KB
/
Day07.hs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
module AOC2017.Day07 (day07a, day07b) where
import AOC2017.Types (Challenge)
import Control.Applicative ((<|>))
import Data.Char (isAlpha)
import Data.Foldable (toList)
import Data.List (sortOn)
import Data.Maybe (mapMaybe, listToMaybe, fromJust)
import Data.Tree
import qualified Data.Map as M
import qualified Data.Set as S
parseLine :: String -> (String, (Int, S.Set String))
parseLine (words->p:w:ws) =
(p, (read w, S.fromList (filter isAlpha <$> drop 1 ws)))
parseLine _ = error "No parse"
-- | Returns the root label and the tree
buildTree
:: M.Map String (Int, S.Set String)
-> (String, Tree Int)
buildTree m = (root, result)
where
allChildren :: S.Set String
allChildren = S.unions (snd <$> toList m)
root :: String
root = S.findMax $ M.keysSet m `S.difference` allChildren
result :: Tree Int
result = flip unfoldTree root $ \p ->
let (w, cs) = m M.! p
in (w, toList cs)
-- | Check if any children are bad; otherwise, check yourself
findBad :: Tree Int -> Maybe Int
findBad t0 = listToMaybe badChildren <|> anomaly
where
badChildren :: [Int]
badChildren = mapMaybe findBad $ subForest t0
weightMap :: M.Map Int [Int]
weightMap = M.fromListWith (++)
. map (\t -> (sum t, [rootLabel t]))
$ subForest t0
anomaly :: Maybe Int
anomaly = case sortOn (length . snd) (M.toList weightMap) of
-- end of the line
[] -> Nothing
-- all weights match
[_] -> Nothing
-- exactly one anomaly
[(wTot1, [w]),(wTot2,_)] -> Just (w + (wTot2 - wTot1))
-- should not happen
_ -> error "More than one anomaly for node"
parse :: String -> (String, Tree Int)
parse = buildTree . M.fromList . map parseLine . lines
day07a :: Challenge
day07a = fst . parse
day07b :: Challenge
day07b = show . fromJust . findBad . snd . parse