Quick actions

cmd+k|ctrl+k

Navigation

Languages

aoc2017-07-2

Snippet info

Language

Python

Visibility

public

Author

christophe.guillon.perso

Created

2017-12-09T02:21:04Z

Updated

2017-12-09T02:21:04Z

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)
INFO