Solve 2019 day 9

2019/d9.py

@@ -0,0 +1,153 @@
+from lib import get_data, str_to_ints
+
+data = "109,1,204,-1,1001,100,1,100,1008,100,16,101,1006,101,0,99"
+# data = "104,1125899906842624,99"
+# data = "1102,34915192,34915192,7,4,7,99,0"
+
+
+class Amp:
+    def __init__(self, xs):
+        self.xs = list(xs)
+        self.xs += [0 for _ in range(1000)]
+        self.i = 0
+        self.inputs = []
+        self.outputs = []
+        self.done = False
+        self.rel_base = 0
+
+    def feed(self, input):
+        self.inputs.append(input)
+
+    def pop(self):
+        v = self.outputs[0]
+        self.outputs = self.outputs[1:]
+        return v
+
+    def get_param(self, offset, mode):
+        if mode == 0:
+            p = self.xs[self.xs[offset]]
+        elif mode == 1:
+            p = self.xs[offset]
+        elif mode == 2:
+            assert self.rel_base + offset >= 0
+            p = self.xs[self.rel_base + self.xs[offset]]
+        else:
+            assert False
+        return p
+
+    def get_addr(self, offset, mode):
+        if mode == 0:
+            return self.xs[offset]
+        elif mode == 2:
+            return self.rel_base + self.xs[offset]
+        else:
+            assert False
+
+    def go_all(self):
+        while not self.done:
+            self.go()
+            while self.outputs:
+                print(self.pop(), end=",")
+
+    def go(self):
+        xs = self.xs
+        i = self.i
+        while i < len(xs):
+            assert xs[i] >= 0
+            inst = str(xs[i])
+            inst = "0" * (5 - len(inst)) + inst
+            assert len(inst) == 5
+            op = int(inst[3:5])
+            mode_p1 = int(inst[2])
+            mode_p2 = int(inst[1])
+            mode_p3 = int(inst[0])
+            match op:
+                case 1:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    p2 = self.get_param(i + 2, mode_p2)
+                    addr = self.get_addr(i + 3, mode_p3)
+                    xs[addr] = p1 + p2
+                    i += 4
+                case 2:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    p2 = self.get_param(i + 2, mode_p2)
+                    addr = self.get_addr(i + 3, mode_p3)
+                    xs[addr] = p1 * p2
+                    i += 4
+                case 3:
+                    # read input
+                    assert len(self.inputs) > 0
+                    addr = self.get_addr(i + 1, mode_p1)
+                    xs[addr] = self.inputs[0]
+                    self.inputs = self.inputs[1:]
+                    i += 2
+                case 4:
+                    # output
+                    v = self.get_param(i + 1, mode_p1)
+                    self.outputs.append(v)
+                    i += 2
+                    self.i = i
+                    return
+                case 99:
+                    self.done = True
+                    return
+                case 5:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    p2 = self.get_param(i + 2, mode_p2)
+                    if p1 != 0:
+                        i = p2
+                    else:
+                        i += 3
+                case 6:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    p2 = self.get_param(i + 2, mode_p2)
+                    if p1 == 0:
+                        i = p2
+                    else:
+                        i += 3
+                case 7:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    p2 = self.get_param(i + 2, mode_p2)
+                    addr = self.get_addr(i + 3, mode_p3)
+                    if p1 < p2:
+                        xs[addr] = 1
+                    else:
+                        xs[addr] = 0
+                    i += 4
+                case 8:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    p2 = self.get_param(i + 2, mode_p2)
+                    addr = self.get_addr(i + 3, mode_p3)
+                    if p1 == p2:
+                        xs[addr] = 1
+                    else:
+                        xs[addr] = 0
+                    i += 4
+                case 9:
+                    p1 = self.get_param(i + 1, mode_p1)
+                    self.rel_base += p1
+                    i += 2
+                case _:
+                    assert False
+            self.i = i
+
+
+def part_1(data):
+    xs_orig = str_to_ints(data)
+
+    a = Amp(xs_orig)
+    a.feed(1)
+    a.go_all()
+
+    a = Amp(xs_orig)
+    a.feed(2)
+    a.go_all()
+
+
+def main():
+    data = get_data(__file__)
+    part_1(data)
+
+
+if __name__ == "__main__":
+    main()