Well fuck the entire file with tests and runninig won't fit in the reply so here's the relevant part, any idea how to tidy it up?
from util import to_string_rows, get_input_string_rows, test, run
from itertools import combinations
from re import compile
def main():
input = get_input_string_rows(7)
RE_DISC = compile(r'(\S+) \((\d+)\)( -> ([\S ]+))?')
class Disc(object):
def __init__(self, name, weight, children):
self.name = name
self.weight = int(weight)
self.children = [] if children is None else children.strip().split(', ')
self.parent = None
def get_children(discs, disc):
return [discs.get(child) for child in disc.children]
def get_weight(discs, disc):
return disc.weight + sum([get_weight(discs, child) for child in get_children(discs, disc)])
def get_root_name(discs, root):
return root.name
def get_weight_correction(discs, root):
prev = dict()
while True:
weights = {child.name: get_weight(discs, child) for child in get_children(discs, root)}
values = weights.values()
diff = max(values) - min(values)
if diff == 0:
return root.weight - (max(prev.values()) - min(prev.values()))
prev = weights
root = discs.get(list(weights.keys())[list(values).index(max(values))])
return 0
def solve(input, fun):
discs = {
m.group(1): Disc(m.group(1), m.group(2), m.group(4))
for m in [RE_DISC.match(row) for row in input]
}
for disc in discs.values():
for child in get_children(discs, disc):
child.parent = disc
root = list(filter(lambda disc: not disc.parent, discs.values()))[0]
return fun(discs, root)
