python(day20): Optimize and tidy up

python/day20.py

@@ -2,15 +2,36 @@ import fileinput
 from pprint import pprint
 import itertools
 from copy import deepcopy
+from typing import List, Tuple, Optional, Dict
 
 N = 10
+SNEK_POSITIONS = (
+    (0, 0),
+    (1, 1),
+    (4, 1),
+    (5, 0),
+    (6, 0),
+    (7, 1),
+    (10, 1),
+    (11, 0),
+    (12, 0),
+    (13, 1),
+    (16, 1),
+    (17, 0),
+    (18, 0),
+    (18, -1),
+    (19, 0),
+)
 
-#
+Tile = List[List[bool]]
-# Parse the input
+Position = Tuple[int, int]
-#
+Extremes = Tuple[int, int, int, int]
-current_tile = None
+
-tiles = []
+
-for line in fileinput.input():
+def parse() -> List[Tuple[int, Tile]]:
+    current_tile = None
+    tiles = []
+    for line in fileinput.input():
         line = line.strip()
         if not line:
             continue
@@ -23,222 +44,169 @@ for line in fileinput.input():
             current_tile = []
         else:
             assert len(line) == N
-        current_tile.append(list(line))
+            current_tile.append([c == "#" for c in line])
 
-assert len(current_tile) == N
+    assert len(current_tile) == N
-tiles.append((current_id, current_tile))
+    tiles.append((current_id, current_tile))
+    return tiles
 
 
-def aligns_right(left, right):
+def aligns_right(left: Tile, right: Tile) -> bool:
-    return all(left[y][-1] == right[y][0] for y in range(len(left)))
+    return all(left_row[-1] == right_row[0] for (left_row, right_row) in zip(left, right))
 
 
-def aligns_bottom(top, bottom):
+def aligns_bottom(top: Tile, bottom: Tile) -> bool:
-    return all(top[-1][x] == bottom[0][x] for x in range(len(top)))
+    return top[-1] == bottom[0]
 
 
-def aligns(apos, b):
+def aligns(a: Tile, b: Tile) -> Optional[Position]:
-    ax, ay = apos
-    _, a = kek[(ax, ay)]
-    if aligns_right(a, b):
-        return 1, 0
-    if aligns_right(b, a):
-        return -1, 0
     if aligns_bottom(a, b):
         return 0, 1
     if aligns_bottom(b, a):
         return 0, -1
+    if aligns_right(a, b):
+        return 1, 0
+    if aligns_right(b, a):
+        return -1, 0
     return None
 
 
-def rotate(a):
+def rotate(tile: Tile) -> Tile:
-    output = deepcopy(a)
+    output = list(reversed(tile))
-    for y in range(len(a)):
+    for y in range(len(tile)):
-        for x in range(len(a)):
+        for x in range(y):
-            output[-1 - x][y] = a[y][x]
+            output[y][x], output[x][y] = output[x][y], output[y][x]
     return output
 
 
-def flip(a):
+def flip(a: Tile) -> Tile:
-    output = [None] * len(a)
+    return list(reversed(a))
-    for y in range(len(a)):
-        output[y] = list(reversed(a[y]))
-    return output
 
 
-def rotate_align(a_pos, b):
+def rotate_align(a: Tile, b: Tile) -> Optional[Tuple[Tile, Position]]:
     """
     Rotates and flips a and checks if it aligns for every possible orientation.
     """
-    if d := aligns(a_pos, b):
+    for _ in range(4):
-        return b, d
+        if pos := aligns(a, b):
+            return b, pos
         bf = flip(b)
-    if d := aligns(a_pos, bf):
-        return bf, d
 
+        if pos := aligns(a, bf):
+            return bf, pos
         b = rotate(b)
-    if d := aligns(a_pos, b):
+    return None
-        return b, d
-    bf = flip(b)
-    if d := aligns(a_pos, bf):
-        return bf, d
-
-    b = rotate(b)
-    if d := aligns(a_pos, b):
-        return b, d
-    bf = flip(b)
-    if d := aligns(a_pos, bf):
-        return bf, d
-
-    b = rotate(b)
-    if d := aligns(a_pos, b):
-        return b, d
-    bf = flip(b)
-    if d := aligns(a_pos, bf):
-        return bf, d
-
-    return None, None
 
 
-positions = {tiles[0][0]: ((0, 0), tiles[0][1])}
+def part1(
-kek = {(0, 0): tiles[0]}
+    tiles: List[Tuple[int, Tile]]
+) -> Tuple[Dict[Position, Tuple[int, Tile]], Extremes]:
+    tile_positions = {tiles[0][0]: (0, 0)}
+    position_tiles = {(0, 0): tiles[0]}
 
-while len(positions) != len(tiles):
+    while len(tile_positions) != len(tiles):
         for a_id, _ in tiles:
-        ((a_x, a_y), a_tile) = positions.get(a_id, ((None, None), None))
+            try:
-        if not a_tile:
+                (a_x, a_y) = a_pos = tile_positions[a_id]
+                _, a_tile = position_tiles[a_pos]
+            except KeyError:
                 continue
 
             for b_id, b_tile in tiles:
-            if b_id in positions or a_id == b_id:
+                if b_id in tile_positions or a_id == b_id:
                     continue
 
-            transformed, dpos = rotate_align((a_x, a_y), b_tile)
+                aligned = rotate_align(a_tile, b_tile)
-            if dpos is not None:
+                if aligned is not None:
-                dx, dy = dpos
+                    transformed, b_pos = aligned
+                    dx, dy = b_pos
                     b_x = a_x + dx
                     b_y = a_y + dy
-                positions[b_id] = (b_x, b_y), transformed
+                    tile_positions[b_id] = (b_x, b_y)
-                kek[(b_x, b_y)] = b_id, transformed
+                    position_tiles[(b_x, b_y)] = b_id, transformed
 
-min_y = min(y for ((_, y), _) in positions.values())
+    min_y = min(y for (_, y) in tile_positions.values())
-max_y = max(y for ((_, y), _) in positions.values())
+    max_y = max(y for (_, y) in tile_positions.values())
-min_x = min(x for ((x, _), _) in positions.values())
+    min_x = min(x for (x, _) in tile_positions.values())
-max_x = max(x for ((x, _), _) in positions.values())
+    max_x = max(x for (x, _) in tile_positions.values())
 
-tilemap = {pos: (id, tile) for id, (pos, tile) in positions.items()}
+    bl, _ = position_tiles[(min_x, min_y)]
-bl, _ = tilemap[(min_x, min_y)]
+    br, _ = position_tiles[(max_x, min_y)]
-br, _ = tilemap[(max_x, min_y)]
+    tl, _ = position_tiles[(min_x, max_y)]
-tl, _ = tilemap[(min_x, max_y)]
+    tr, _ = position_tiles[(max_x, max_y)]
-tr, _ = tilemap[(max_x, max_y)]
+    print("Part 1:", tl * tr * bl * br)
-print("Part 1:", tl * tr * bl * br)
+
+    return position_tiles, (min_x, max_x, min_y, max_y)
 
 
-def is_snek(b, x, y):
+def is_snek(image: Tile, start_x: int, start_y: int) -> bool:
-    for (x_, y_) in (
+    for (x, y) in SNEK_POSITIONS:
-        (0, 0),
-        (1, 1),
-        (4, 1),
-        (5, 0),
-        (6, 0),
-        (7, 1),
-        (10, 1),
-        (11, 0),
-        (12, 0),
-        (13, 1),
-        (16, 1),
-        (17, 0),
-        (18, 0),
-        (18, -1),
-        (19, 0),
-    ):
         try:
-            if b[y_ + y][x_ + x] == ".":
+            if not image[y + start_y][x + start_x]:
                 return False
         except:
             return False
     return True
 
 
-def find_sneks(b):
+def find_sneks(b: Tile) -> List[Position]:
     sneks = []
     for (x, y) in itertools.product(range(8 * 12), repeat=2):
         if is_snek(b, x, y):
-            print(x, y, "is snek")
             sneks.append((x, y))
     return sneks
 
 
-def rotate_find_sneks(b):
+def rotate_find_sneks(b: Tile) -> Tuple[Tile, List[Position]]:
-    if d := find_sneks(b):
+    for _ in range(4):
-        return b, d
+        if position := find_sneks(b):
+            return b, position
         bf = flip(b)
-    if d := find_sneks(bf):
-        return bf, d
 
+        if position := find_sneks(bf):
+            return bf, position
         b = rotate(b)
-    if d := find_sneks(b):
+    raise RuntimeError("no sneks found")
-        return b, d
-    bf = flip(b)
-    if d := find_sneks(bf):
-        return bf, d
-
-    b = rotate(b)
-    if d := find_sneks(b):
-        return b, d
-    bf = flip(b)
-    if d := find_sneks(bf):
-        return bf, d
-
-    b = rotate(b)
-    if d := find_sneks(b):
-        return b, d
-    bf = flip(b)
-    if d := find_sneks(bf):
-        return bf, d
 
 
-def remove_snek(snek):
+def remove_snek(image: Tile, snek_position: Position) -> None:
     # :(
-    (x, y) = snek
+    (x, y) = snek_position
-    for (x_, y_) in (
+    for (x_, y_) in SNEK_POSITIONS:
-        (0, 0),
+        image[y + y_][x + x_] = False
-        (1, 1),
-        (4, 1),
-        (5, 0),
-        (6, 0),
-        (7, 1),
-        (10, 1),
-        (11, 0),
-        (12, 0),
-        (13, 1),
-        (16, 1),
-        (17, 0),
-        (18, 0),
-        (18, -1),
-        (19, 0),
-    ):
-        big_chungus[y + y_][x + x_] = '.'
 
 
-chungus = [["."] * 8 * (max_x - min_x + 1) for _ in range(8 * (max_y - min_y + 1))]
+def part2(
-for (x, y), (id, tile) in tilemap.items():
+    position_tiles: Dict[Position, Tuple[int, Tile]],
+    extremes: Tuple[int, int, int, int],
+) -> None:
+    min_x, max_x, min_y, max_y = extremes
+    image = [[False] * 8 * (max_x - min_x + 1) for _ in range(8 * (max_y - min_y + 1))]
+    for (x, y), (_, tile) in position_tiles.items():
         x_ = (x - min_x) * 8
         y_ = (y - min_y) * 8
         for i, row in enumerate(tile[1:-1]):
             for j, c in enumerate(row[1:-1]):
-            chungus[i + y_][j + x_] = c
+                image[i + y_][j + x_] = c
 
-# for y, row in enumerate(chungus):
+    rotated_image, snek_positions = rotate_find_sneks(image)
-#     print("".join(row))
+    for snek_position in snek_positions:
+        remove_snek(rotated_image, snek_position)
 
-big_chungus, sneks = rotate_find_sneks(chungus)
+    part2 = 0
-for snek in sneks:
+    for row in rotated_image:
-    remove_snek(snek)
-
-part2 = 0
-for row in big_chungus:
         for c in row:
-        if c == "#":
+            if c:
                 part2 += 1
-print("Part 2:", part2)
+
+    print("Part 2:", part2)
+
+
+def main() -> None:
+    tiles = parse()
+    position_tiles, extremes = part1(tiles)
+    part2(position_tiles, extremes)
+
+
+if __name__ == "__main__":
+    main()