Add 2022 solutions

2022/d21.py

@@ -0,0 +1,105 @@
+from lib import *
+
+EXAMPLE = """root: pppw + sjmn
+dbpl: 5
+cczh: sllz + lgvd
+zczc: 2
+ptdq: humn - dvpt
+dvpt: 3
+lfqf: 4
+humn: 5
+ljgn: 2
+sjmn: drzm * dbpl
+sllz: 4
+pppw: cczh / lfqf
+lgvd: ljgn * ptdq
+drzm: hmdt - zczc
+hmdt: 32
+"""
+
+def resolve(monkeys, s):
+    equ = monkeys[s]
+    try:
+        return int(equ)
+    except ValueError:
+        pass
+    l, op, r = equ.split(" ")
+    l = resolve(monkeys, l)
+    r = resolve(monkeys, r)
+    res = eval(f"{l}{op}{r}")
+    monkeys[s] = res
+    return res
+
+
+def resolve2(monkeys, s):
+    if s == "humn":
+        return tuple()
+
+    equ = monkeys[s]
+    try:
+        return int(equ)
+    except ValueError:
+        pass
+
+    l, op, r = equ.split(" ")
+    if s == "root":
+        op = "=="
+
+    l = resolve2(monkeys, l)
+    r = resolve2(monkeys, r)
+
+    if isinstance(l, tuple) or isinstance(r, tuple):
+        return (l, op, r)
+    if op == "/":
+        op = "//"
+    res = eval(f"{l}{op}{r}")
+    monkeys[s] = res
+    return res
+
+
+def simplify(equ, exp):
+    if equ == tuple():
+        return exp
+    l, op, r = equ
+    if op == '/' and isinstance(r, int):
+        return simplify(l, exp * r)
+    elif op == '+' and isinstance(l, int):
+        return simplify(r, exp - l)
+    elif op == '+' and isinstance(r, int):
+        return simplify(l, exp - r)
+    elif op == '-' and isinstance(l, int):
+        return simplify(r, -(exp - l))
+    elif op == '-' and isinstance(r, int):
+        return simplify(l, exp + r)
+    elif op == '*' and isinstance(r, int):
+        return simplify(l, exp // r)
+    elif op == '*' and isinstance(l, int):
+        return simplify(r, exp // l)
+
+
+def solve(input: Input, second=False):
+    res = 0
+    monkeys = {}
+    for line in input.lines():
+        l, r = line.split(": ")
+        monkeys[l] = r
+
+    if not second:
+        return int(resolve(monkeys, "root"))
+    else:
+        equ = resolve2(monkeys, "root")
+        assert equ[1] == '=='
+        return simplify(equ[0], equ[2])
+
+
+def main():
+    DAY_INPUT = "d21.txt"
+    print("Example  1:", solve(Input(EXAMPLE)))
+    print("Solution 1:", solve(Input(DAY_INPUT)))
+    print("Example  2:", solve(Input(EXAMPLE), True))
+    print("Solution 2:", solve(Input(DAY_INPUT), True))
+    assert solve(Input(DAY_INPUT), True) == 3243420789721
+
+
+if __name__ == "__main__":
+    main()