71 lines
2 KiB
Python
71 lines
2 KiB
Python
from collections import namedtuple
|
|
|
|
Node = namedtuple('Node', ['name', 'disc', 'children'])
|
|
|
|
with open('input/day7') as f:
|
|
puzzle_input = []
|
|
for line in f.read().splitlines():
|
|
children = []
|
|
if '->' in line:
|
|
line, rest = line.split('-> ')
|
|
children = rest.split(', ')
|
|
name, weight = line.split()
|
|
weight = int(weight[1:-1])
|
|
|
|
puzzle_input.append(Node(name, weight, tuple(children)))
|
|
|
|
def find_root(nodes):
|
|
names = {node.name for node in nodes}
|
|
supported = {child for node in nodes for child in node.children}
|
|
return (names - supported).pop()
|
|
|
|
print(find_root(puzzle_input))
|
|
|
|
def memoized(f):
|
|
cache = {}
|
|
def _(*args):
|
|
args = tuple(args)
|
|
if args in cache:
|
|
return cache[args]
|
|
val = f(*args)
|
|
cache[args] = val
|
|
return val
|
|
return _
|
|
|
|
def identical(iterable):
|
|
return not iterable or len(set(iterable)) == 1
|
|
|
|
def counts(iterable):
|
|
count = {}
|
|
for i in iterable:
|
|
count[i] = count.get(i, 0) + 1
|
|
return count
|
|
|
|
def unbalanced(nodes):
|
|
by_name = {node.name: node for node in nodes}
|
|
root = by_name[find_root(nodes)]
|
|
|
|
@memoized
|
|
def weight(n):
|
|
return n.disc + sum(weight(by_name[c]) for c in n.children)
|
|
|
|
@memoized
|
|
def is_balanced(n):
|
|
return identical(weight(by_name[c]) for c in n.children)
|
|
|
|
def walk(tree):
|
|
if not is_balanced(tree):
|
|
weights = [weight(by_name[c]) for c in tree.children]
|
|
yield tree.name
|
|
for n in (by_name[c] for c in tree.children):
|
|
yield from walk(n)
|
|
|
|
unbalanced = list(walk(root))[-1]
|
|
children = [by_name[c] for c in by_name[unbalanced].children]
|
|
weights = [weight(c) for c in children]
|
|
weight_counts = counts(weights)
|
|
odd_one = min(children, key=lambda n: weight_counts[weight(n)])
|
|
majority = max(set(weights), key=lambda w: weight_counts[w])
|
|
return majority - sum(weight(by_name[c]) for c in odd_one.children)
|
|
|
|
print(unbalanced(puzzle_input))
|