import Control.Monad import qualified Data.Map.Strict as Map import qualified Data.Map.Lazy as LMap import Data.Map.Strict ((!)) import Data.Maybe import Data.List import qualified Data.Set as Set data Program = Program {progName :: String, progWeight :: Int, progChildren :: [String]} deriving Show parse :: String -> Program parse s = go (words s) where go [name, weight] = Program name (readweight weight) [] go (name : weight : "->" : ch) = Program name (readweight weight) (map (filter (/= ',')) ch) readweight = read . init . tail findParents :: [Program] -> (Map.Map String String, Set.Set String) findParents allprogs = go allprogs Map.empty (Set.fromList (map progName allprogs)) where go [] parmap nopar = (parmap, nopar) go (Program name _ ch : progs) parmap nopar = let parmap' = foldl (\m c -> Map.insert c name m) parmap ch nopar' = foldl (\s c -> Set.delete c s) nopar ch in go progs parmap' nopar' towerWeights :: [Program] -> Map.Map String Int towerWeights progs = resmap where resmap = LMap.fromList [(progName p, go p) | p <- progs] go prog@(Program _ w ch) = w + sum [resmap ! c | c <- ch] alleq :: Eq a => [a] -> Bool alleq (a:b:cs) = a == b && alleq (b:cs) alleq _ = True oneout :: Eq a => [a] -> (Int {-idx-}, a {-mode-}) oneout l = go l 0 where go (a:b:c:_) 0 | a /= b && b == c = (0, b) go [a,b,c] i | a == b && b /= c = (i + 2, a) go (a:b:c:_) i | a /= b && b /= c = (succ i, a) go (_:xs) i = go xs (succ i) main :: IO () main = do input <- liftM lines (readFile "7.in") let progs = map parse input (parmap, nopar) = findParents progs [root] = Set.toList nopar putStrLn root let progmap = Map.fromList [(progName p, p) | p <- progs] tw = towerWeights progs unbalnames = map progName $ filter (\(Program _ _ ch) -> not $ alleq $ map (tw !) ch) progs [unbalname] = unbalnames \\ catMaybes (map (flip Map.lookup parmap) unbalnames) unbal = progmap ! unbalname (wrongidx, otherw) = oneout (map (tw !) (progChildren unbal)) wrongprog = progmap ! (progChildren unbal !! wrongidx) print $ progWeight wrongprog + otherw - (tw ! progName wrongprog)