Optimize day 20 solution

day-20/day-20.py

@@ -3,55 +3,59 @@ from pathlib import Path
 
 
 def enhance(algorithm, in_image, boundary):
-    out_image = []
-    height = len(in_image)
-    width = len(in_image[0])
-    for i in range(-boundary, height + boundary):
-        line = ''
-        valid = False
-        for j in range(-boundary, width + boundary):
-            idx = ''
-            for k in range(i - 1, i + 2):
-                for l in range(j - 1, j + 2):
-                    if 0 <= k < height and \
-                       0 <= l < width and \
-                       '#' == in_image[k][l]:
-                        idx += '1'
-                    else:
-                        idx += '0'
-            idx = int(idx, 2)
-            if not valid and '#' == algorithm[idx]:
-                valid = True
-            line += algorithm[idx]
-        if valid:
-            out_image.append(line)
+    out_image = set()
+    min_x = min([x for x, _ in in_image])
+    max_x = max([x + 1 for x, _ in in_image])
+    min_y = min([y for _, y in in_image])
+    max_y = max([y + 1 for _, y in in_image])
+    for y in range(min_y - boundary, max_y + boundary):
+        for x in range(min_x - boundary, max_x + boundary):
+            idx = 0
+            bit = 8
+            for dy in [-1, 0, 1]:
+                for dx in [-1, 0, 1]:
+                    idx <<= 1
+                    if (x + dx, y + dy) in in_image:
+                        idx |= 1
+            if '#' == algorithm[idx]:
+                out_image.add((x, y))
     return out_image
 
 
+def show(image):
+    min_x = min([x for x, _ in image])
+    max_x = max([x + 1 for x, _ in image])
+    min_y = min([y for _, y in image])
+    max_y = max([y + 1 for _, y in image])
+    for y in range(min_y, max_y):
+        for x in range(min_x, max_x):
+            if (x, y) in image:
+                print('#', end='')
+            else:
+                print('.', end='')
+        print()
+
+
 def part_1(input):
     result = 0
     algorithm = input[0].rstrip()
-    image = []
-    for line in input[2:]:
-        image.append(line.rstrip())
+    image = set([(x, y) for y, line in enumerate(input[2:])
+                 for x, c in enumerate(line) if c == '#'])
     image = enhance(algorithm, image, 3)
     image = enhance(algorithm, image, -1)
-    for line in image:
-        result += line.count('#')
+    result = len(image)
     print("Part 1 result:", result)
 
 
 def part_2(input):
     result = 0
     algorithm = input[0].rstrip()
-    image = []
-    for line in input[2:]:
-        image.append(line.rstrip())
+    image = set([(x, y) for y, line in enumerate(input[2:])
+                 for x, c in enumerate(line) if c == '#'])
     for _ in range(25):
         image = enhance(algorithm, image, 3)
         image = enhance(algorithm, image, -1)
-    for line in image:
-        result += line.count('#')
+    result = len(image)
     print("Part 2 result:", result)