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
63
64
|
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)
|
