Day07.hs

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