Solve 2018 day 16 and add function to download input

2018/d16.py

@@ -0,0 +1,147 @@
+from lib import str_to_ints, get_data
+
+
+def addr(regs, a, b, c):
+    regs[c] = regs[a] + regs[b]
+
+
+def addi(regs, a, b, c):
+    regs[c] = regs[a] + b
+
+
+def mulr(regs, a, b, c):
+    regs[c] = regs[a] * regs[b]
+
+
+def muli(regs, a, b, c):
+    regs[c] = regs[a] * b
+
+
+def banr(regs, a, b, c):
+    regs[c] = regs[a] & regs[b]
+
+
+def bani(regs, a, b, c):
+    regs[c] = regs[a] & b
+
+
+def borr(regs, a, b, c):
+    regs[c] = regs[a] | regs[b]
+
+
+def bori(regs, a, b, c):
+    regs[c] = regs[a] | b
+
+
+def setr(regs, a, _, c):
+    regs[c] = regs[a]
+
+
+def seti(regs, a, _, c):
+    regs[c] = a
+
+
+def gtir(regs, a, b, c):
+    regs[c] = 1 if a > regs[b] else 0
+
+
+def gtri(regs, a, b, c):
+    regs[c] = 1 if regs[a] > b else 0
+
+
+def gtrr(regs, a, b, c):
+    regs[c] = 1 if regs[a] > regs[b] else 0
+
+
+def eqir(regs, a, b, c):
+    regs[c] = 1 if a == regs[b] else 0
+
+
+def eqri(regs, a, b, c):
+    regs[c] = 1 if regs[a] == b else 0
+
+
+def eqrr(regs, a, b, c):
+    regs[c] = 1 if regs[a] == regs[b] else 0
+
+
+OPS = [
+    addr,
+    addi,
+    mulr,
+    muli,
+    banr,
+    bani,
+    borr,
+    bori,
+    setr,
+    seti,
+    gtir,
+    gtri,
+    gtrr,
+    eqir,
+    eqri,
+    eqrr,
+]
+
+
+def part_1(data):
+    examples = []
+
+    lines = data.splitlines()
+    test_prog = []
+    last_i = None
+    for i in range(len(lines)):
+        line = lines[i]
+        if line.startswith("Before:"):
+            regs = str_to_ints(line)
+            inst = str_to_ints(lines[i + 1])
+            regs_after = str_to_ints(lines[i + 2])
+            examples.append((regs, inst, regs_after))
+            last_i = i + 2
+
+    assert last_i is not None
+    for line in lines[last_i:]:
+        if line.strip() != "":
+            test_prog.append(str_to_ints(line))
+
+    r = 0
+    for before_orig, inst, after in examples:
+        ops_correct = []
+        for op in OPS:
+            before = list(before_orig)
+            op(before, *inst[1:])
+            if before == after:
+                ops_correct.append(op)
+        if len(ops_correct) >= 3:
+            r += 1
+    print(r)
+
+    code_to_op = {}
+    while len(OPS) > 0:
+        for before_orig, inst, after in examples:
+            ops_correct = []
+            for op in OPS:
+                before = list(before_orig)
+                op(before, *inst[1:])
+                if before == after:
+                    ops_correct.append(op)
+            if len(ops_correct) == 1:
+                code_to_op[inst[0]] = ops_correct[0]
+                OPS.remove(ops_correct[0])
+
+    regs = [0, 0, 0, 0]
+    for line in test_prog:
+        op_code = line[0]
+        vals = line[1:]
+        code_to_op[op_code](regs, *vals)
+    print(regs[0])
+
+
+def main():
+    data = get_data(__file__)
+    part_1(data)
+
+
+if __name__ == "__main__":
+    main()