Day 7

day7.py

@@ -0,0 +1,71 @@
+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))