Solve 2019 day 11

2019/d11.py

@@ -0,0 +1,93 @@
+from lib import get_data, str_to_ints, add2
+from d9 import Amp
+
+DIRS = [(-1, 0), (0, 1), (1, 0), (0, -1)]
+BLACK = 0
+WHITE = 1
+
+
+def part_1(data):
+    xs = str_to_ints(data)
+    a = Amp(xs)
+    current_dir_idx = 0
+    pos = (0, 0)
+    white_panels = set()
+    got_painted = set()
+
+    while not a.done:
+        if pos in white_panels:
+            a.feed(WHITE)
+        else:
+            a.feed(BLACK)
+        a.go()
+        a.go()
+        while a.outputs:
+            paint_color = a.pop()
+            if paint_color == BLACK:
+                white_panels.remove(pos)
+            elif paint_color == WHITE:
+                white_panels.add(pos)
+            else:
+                assert False
+            got_painted.add(pos)
+            turn_dir_code = a.pop()
+            turn_dir = -1 if turn_dir_code == 0 else 1
+            current_dir_idx = (current_dir_idx + turn_dir) % len(DIRS)
+            pos = add2(pos, DIRS[current_dir_idx])
+    print(len(got_painted))
+
+
+def part_2(data):
+    xs = str_to_ints(data)
+    a = Amp(xs)
+    current_dir_idx = 0
+    pos = (0, 0)
+    white_panels = set([pos])
+    got_painted = set()
+
+    while not a.done:
+        if pos in white_panels:
+            a.feed(WHITE)
+        else:
+            a.feed(BLACK)
+        a.go()
+        a.go()
+        while a.outputs:
+            paint_color = a.pop()
+            if paint_color == BLACK:
+                if pos in white_panels:
+                    white_panels.remove(pos)
+            elif paint_color == WHITE:
+                white_panels.add(pos)
+            else:
+                assert False
+            got_painted.add(pos)
+            turn_dir_code = a.pop()
+            turn_dir = -1 if turn_dir_code == 0 else 1
+            current_dir_idx = (current_dir_idx + turn_dir) % len(DIRS)
+            pos = add2(pos, DIRS[current_dir_idx])
+
+    white_panels = [(p[1], p[0]) for p in white_panels]
+    xs = [p[0] for p in white_panels]
+    ys = [p[1] for p in white_panels]
+    white_panels_set = set(white_panels)
+    x_min, x_max = min(xs), max(xs)
+    y_min, y_max = min(ys), max(ys)
+    for y in range(y_min, y_max + 1):
+        s = ""
+        for x in range(x_min, x_max + 1):
+            if (x, y) in white_panels_set:
+                s += "X"
+            else:
+                s += " "
+        print(s)
+
+
+def main():
+    data = get_data(__file__)
+    part_1(data)
+    part_2(data)
+
+
+if __name__ == "__main__":
+    main()