aoc2017-07-2
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051
import sys
import re
input = sys.stdin.read()
lines = input.split("\n")
msuccs = {}
mwei = {}
for line in lines:
line = re.sub("[-(),><]", "", line)
w = line.split()
src = w[0]
wei = int(w[1])
succs = w[2:]
msuccs[src] = succs
mwei[src] = wei
parents = {}
for root in msuccs:
for n in msuccs[root]:
parents[n] = root
root = list(filter(lambda x: x not in parents, msuccs))[0]
def dobalance(root,msuccs, wei):
dfs = []
twei = {}
def dodfs(root):
for n in msuccs[root]:
ret = dodfs(n)
if ret: return ret
weights = list(map(lambda x: twei[x], msuccs[root]))
setw = set(weights)
if len(setw) > 1:
for k in setw:
if weights.count(k) == 1:
i = weights.index(k)
j = (i+1)%len(weights)
diff = weights[i] - weights[j]
nwei = wei[msuccs[root][i]] - diff
return (root, nwei)
twei[root] = wei[root]
for n in msuccs[root]:
twei[root] += twei[n]
#print(root, wei[root], twei[root], weights)
dfs.append(root)
return dodfs(root)
balance = dobalance(root, msuccs, mwei)
print(balance)
Python
INFO