Refactor 2024 day 6 to be a little bit faster.

src/holt59/aoc/2024/day5.py

@@ -9,7 +9,7 @@ def in_correct_order(update: list[int], requirements: dict[int, set[int]]) -> bo
     )
 
 
-def fix_order(
+def to_correct_order(
     update: list[int],
     requirements: dict[int, set[int]],
     max_update_length: int | None = None,
@@ -39,7 +39,7 @@ def fix_order(
     return update
 
 
-part1, part2 = sys.stdin.read().split("\n\n")
+part1, part2 = sys.stdin.read().strip().split("\n\n")
 
 requirements: dict[int, set[int]] = defaultdict(set)
 for line in part1.splitlines():
@@ -55,7 +55,7 @@ answer_1 = sum(
 )
 
 answer_2 = sum(
-    fix_order(update, requirements, len(update) // 2 + 1)[-1]
+    to_correct_order(update, requirements, len(update) // 2 + 1)[-1]
     for update in updates
     if not in_correct_order(update, requirements)
 )

src/holt59/aoc/2024/day6.py

@@ -1,29 +1,94 @@
+import itertools as it
 import sys
 from typing import TypeAlias
 
+NodeType: TypeAlias = tuple[tuple[int, int], tuple[int, int]]
+EdgesType: TypeAlias = dict[NodeType, tuple[NodeType, set[tuple[int, int]]]]
+
+ROTATE = {(-1, 0): (0, 1), (0, 1): (1, 0), (1, 0): (0, -1), (0, -1): (-1, 0)}
+
+START_NODE: NodeType = ((-2, -2), (-1, 0))
+FINAL_POS: tuple[int, int] = (-1, -1)
+
 
 def move(
     lines: list[str], pos: tuple[int, int], dir: tuple[int, int]
-) -> tuple[tuple[int, int] | None, list[tuple[int, int]]]:
+) -> tuple[tuple[int, int] | None, set[tuple[int, int]]]:
     n_rows, n_cols = len(lines), len(lines[0])
     row, col = pos
-    marked: list[tuple[int, int]] = []
+
+    marked: set[tuple[int, int]] = set()
+    final_pos: tuple[int, int] | None = None
 
     while True:
-        marked.append((row, col))
+        marked.add((row, col))
 
         if not (0 <= row + dir[0] < n_rows and 0 <= col + dir[1] < n_cols):
-            pos = None
+            final_pos = None
             break
 
         if lines[row + dir[0]][col + dir[1]] != ".":
-            pos = (row, col)
+            final_pos = (row, col)
             break
 
         row += dir[0]
         col += dir[1]
 
-    return pos, marked
+    return final_pos, marked
+
+
+def compute_graph(lines: list[str], start_node: NodeType):
+    n_rows, n_cols = len(lines), len(lines[0])
+
+    edges: EdgesType = {}
+
+    start_pos, start_dir = start_node
+    end_pos, marked = move(lines, start_pos, start_dir)
+    assert end_pos is not None
+    edges[START_NODE] = ((end_pos, start_dir), marked)
+
+    for row, col in it.product(range(n_rows), range(n_cols)):
+        if lines[row][col] != "#":
+            continue
+
+        for start_pos, start_dir in (
+            ((row - 1, col), (1, 0)),
+            ((row + 1, col), (-1, 0)),
+            ((row, col - 1), (0, 1)),
+            ((row, col + 1), (0, -1)),
+        ):
+            if 0 <= start_pos[0] < n_rows and 0 <= start_pos[1] < n_cols:
+                end_pos, marked = move(lines, start_pos, ROTATE[start_dir])
+
+                edges[start_pos, start_dir] = (
+                    (end_pos or FINAL_POS, ROTATE[start_dir]),
+                    marked,
+                )
+
+    return edges
+
+
+def is_loop(lines: list[str], edges: EdgesType, position: tuple[int, int]):
+    row, col = position
+    current_node = START_NODE
+    found: set[NodeType] = set()
+
+    while current_node[0] != FINAL_POS and current_node not in found:
+        found.add(current_node)
+
+        target_node, edge_marked = edges[current_node]
+
+        if (row, col) in edge_marked:
+            # need to break the edge
+            target_dir = target_node[1]
+            end_pos, _ = move(
+                lines, (row - target_dir[0], col - target_dir[1]), ROTATE[target_dir]
+            )
+            current_node = (end_pos or FINAL_POS, ROTATE[target_dir])
+        else:
+            current_node = target_node
+
+    return current_node in found
 
 
 def print_grid(
@@ -42,65 +107,28 @@ def print_grid(
     print()
 
 
-ROTATE = {(-1, 0): (0, 1), (0, 1): (1, 0), (1, 0): (0, -1), (0, -1): (-1, 0)}
-
-
+# read lines
 lines = sys.stdin.read().splitlines()
-n_rows, n_cols = len(lines), len(lines[0])
+
+# find and delete original position
 start_pos = next(
-    (i, j) for i in range(n_rows) for j in range(n_cols) if lines[i][j] == "^"
+    (i, j) for i, row in enumerate(lines) for j, col in enumerate(row) if col == "^"
 )
 lines[start_pos[0]] = lines[start_pos[0]].replace("^", ".")
 
+# compute edges from the map
+edges = compute_graph(lines, (start_pos, (-1, 0)))
+
 # part 1
 marked: set[tuple[int, int]] = set()
+current_node = START_NODE
 
-current_pos: tuple[int, int] | None = start_pos
-current_dir = (-1, 0)
-
-while current_pos is not None:
-    # print_grid(lines, marked, current_pos)
-    new_pos, new_marked = move(lines, current_pos, current_dir)
-
-    marked = marked.union(new_marked)
-
-    current_pos = new_pos
-    current_dir = ROTATE[current_dir]
-
-# print_grid(lines, marked, current_pos)
+while current_node[0] != FINAL_POS:
+    current_node, current_marked = edges[current_node]
+    marked = marked.union(current_marked)
 
 answer_1 = len(marked)
-
-answer_2 = 0
-for row in range(n_rows):
-    for col in range(n_cols):
-        if (row, col) == start_pos or lines[row][col] != ".":
-            continue
-
-        chars = list(map(list, lines))
-
-        chars[row][col] = "#"
-
-        current_pos = start_pos
-        current_dir = (-1, 0)
-
-        found: set[tuple[tuple[int, int], tuple[int, int]]] = set()
-
-        while current_pos is not None:
-            if current_pos != start_pos:
-                found.add((current_pos, current_dir))
-
-            new_pos, new_marked = move(
-                ["".join(line) for line in chars], current_pos, current_dir
-            )
-
-            current_pos = new_pos
-            current_dir = ROTATE[current_dir]
-
-            if (current_pos, current_dir) in found:
-                print(row, col)
-                answer_2 += 1
-                break
-
 print(f"answer 1 is {answer_1}")
+
+answer_2 = sum(is_loop(lines, edges, pos) for pos in marked if pos != start_pos)
 print(f"answer 2 is {answer_2}")