day 21 python part 2

day21/sol.py

@@ -36,4 +36,68 @@ def eval(x):
     args = [eval(a) for a in x[1:]]
     return optab[op](*args)
 
-print(eval(scope['root']))
+root = scope['root']
+print(eval(root))
+
+def solve(x):
+    _, left, right = x
+
+    humn = scope['humn']
+    path = set()
+
+    def find(x):
+        if x is humn:
+            path.add(x)
+            return True
+        if type(x) is tuple:
+            for a in x:
+                if find(a):
+                    path.add(x)
+                    return True
+            else:
+                return False
+        return False
+
+    x, y = left, right
+    if not find(x):
+        find(y)
+        x, y = y, x
+
+    #print(x)
+    assert x in path
+    y = eval(y)
+
+    while x is not humn:
+        assert len(x) == 3
+        op, a, b = x
+        if op == '/':
+            if a in path:
+                y = ('*', y, b)
+            else:
+                raise ValueError("ruh-roh")
+        elif op == '*':
+            if a in path:
+                y = ('/', y, b)
+            else:
+                y = ('/', y, a)
+        elif op == '+':
+            if a in path:
+                y = ('-', y, b)
+            else:
+                y = ('-', y, a)
+        elif op == '-':
+            if a in path:
+                # x - b = y => x = y + b
+                y = ('+', y, b)
+            else:
+                # a - x = y  => x = a - y
+                y = ('-', a, y)
+        else:
+            raise ValueError("unknown op "+repr(op))
+        if a in path:
+            x = a
+        else:
+            x = b
+    return eval(y)
+
+print(solve(root))