Simplify test input

puzzles/1.py

@@ -1,6 +1,6 @@
 from collections import Counter
 
-test_input_p1 = test_input_p2 = """
+test_input = """
 3   4
 4   3
 2   5

puzzles/2.py

@@ -3,7 +3,7 @@ from typing import Iterator
 
 from more_itertools import ilen
 
-test_input_p1 = test_input_p2 = """
+test_input = """
 7 6 4 2 1
 1 2 7 8 9
 9 7 6 2 1

puzzles/3.py

@@ -1,8 +1,12 @@
 import re
 from functools import reduce
 
-test_input_p1 = "xmul(2,4)%&mul[3,7]!@^do_not_mul(5,5)+mul(32,64]then(mul(11,8)mul(8,5))"
+test_input_p1 = (
-test_input_p2 = "xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))"
+    "xmul(2,4)%&mul[3,7]!@^do_not_mul(5,5)+mul(32,64]then(mul(11,8)mul(8,5))"
+)
+test_input_p2 = (
+    "xmul(2,4)&mul[3,7]!^don't()_mul(5,5)+mul(32,64](mul(11,8)undo()?mul(8,5))"
+)
 
 test_solution_p1 = 161
 test_solution_p2 = 48

puzzles/4.py

@@ -1,7 +1,7 @@
 from typing import Iterator
 
 
-test_input_p1 = test_input_p2 = """
+test_input = """
 MMMSXXMASM
 MSAMXMSMSA
 AMXSXMAAMM
@@ -39,8 +39,7 @@ def _scan_word_grid(
     yield from (row for row in char_grid)
     yield from (tuple(col) for col in zip(*char_grid))
     yield from (
-        _diagonal(char_grid, i)
+        _diagonal(char_grid, i) for i in range(-len(char_grid) + 1, len(char_grid[0]))
-        for i in range(-len(char_grid) + 1, len(char_grid[0]))
     )
     yield from (
         _diagonal(tuple(reversed(char_grid)), i)
@@ -57,28 +56,29 @@ def _diagonal(grid: tuple[tuple[str, ...], ...], offset=0) -> tuple[str, ...]:
 def _count_xmas(line: tuple[str, ...]) -> int:
     return sum(
         1
-        for i
+        for i in range(len(line) - 3)
-        in range(len(line) - 3)
         if line[i : i + 4] in (("X", "M", "A", "S"), ("S", "A", "M", "X"))
     )
 
-ThreeSquare = tuple[
+
-    tuple[str, str, str],
+ThreeSquare = tuple[tuple[str, str, str], tuple[str, str, str], tuple[str, str, str]]
-    tuple[str, str, str],
+
-    tuple[str, str, str]
-]
 
 def _three_squares(word_grid: tuple[tuple[str, ...], ...]) -> Iterator[ThreeSquare]:
     yield from (
-        (word_grid[i][j:j+3], word_grid[i+1][j:j+3], word_grid[i+2][j:j+3])
+        (
+            word_grid[i][j : j + 3],
+            word_grid[i + 1][j : j + 3],
+            word_grid[i + 2][j : j + 3],
+        )
         for i in range(len(word_grid) - 2)
         for j in range(len(word_grid[0]) - 2)
     )
 
+
 def _has_cross_mas(square: ThreeSquare):
     diag_1 = (square[0][0], square[1][1], square[2][2])
     diag_2 = (square[0][2], square[1][1], square[2][0])
-    return (
+    return (diag_1 == ("M", "A", "S") or diag_1 == ("S", "A", "M")) and (
-        (diag_1 == ("M", "A", "S") or diag_1 == ("S", "A", "M"))
+        diag_2 == ("M", "A", "S") or diag_2 == ("S", "A", "M")
-        and (diag_2 == ("M", "A", "S") or diag_2 == ("S", "A", "M"))
+    )
-    )

puzzles/5.py

@@ -0,0 +1,111 @@
+from typing import Iterator
+
+
+test_input = """
+47|53
+97|13
+97|61
+97|47
+75|29
+61|13
+75|53
+29|13
+97|29
+53|29
+61|53
+97|53
+61|29
+47|13
+75|47
+97|75
+47|61
+75|61
+47|29
+75|13
+53|13
+
+75,47,61,53,29
+97,61,53,29,13
+75,29,13
+75,97,47,61,53
+61,13,29
+97,13,75,29,47
+""".strip()
+
+test_solution_p1 = 143
+test_solution_p2 = 123
+
+
+def solve_p1(puzzle_input: str) -> int:
+    rules, updates = _parse_rules_and_updates(puzzle_input)
+    return sum(
+        update[len(update) // 2]
+        for update in updates
+        if _is_correctly_ordered(rules, update)
+    )
+
+
+def solve_p2(puzzle_input: str) -> int:
+    rules, updates = _parse_rules_and_updates(puzzle_input)
+    corrected_updates = (
+        _fix_update(rules, update)
+        for update in updates
+        if not _is_correctly_ordered(rules, update)
+    )
+    return sum(update[len(update) // 2] for update in corrected_updates)
+
+
+OrderingRules = dict[int, set[int]]
+Update = tuple[int, ...]
+
+
+def _parse_rules_and_updates(
+    puzzle_input: str,
+) -> tuple[OrderingRules, Iterator[Update]]:
+    ordering_rules, _, updates = puzzle_input.partition("\n\n")
+
+    rules: OrderingRules = {}
+    for line in ordering_rules.strip().split("\n"):
+        a, _, b = line.partition("|")
+        a, b = int(a), int(b)
+        if b not in rules:
+            rules[b] = set()
+        rules[b].add(a)
+
+    return (
+        rules,
+        (tuple(map(int, line.split(","))) for line in updates.strip().split("\n")),
+    )
+
+
+def _is_correctly_ordered(rules: OrderingRules, update: Update) -> bool:
+    contains = set(update)
+    seen = set()
+
+    for i in update:
+        for j in rules.get(i, ()):
+            if j in contains and j not in seen:
+                return False
+        seen.add(i)
+
+    return True
+
+
+def _fix_update(rules: OrderingRules, update: Update) -> Update:
+    contains = set(update)
+    seen = set()
+
+    fixed_update = list()
+
+    def _fix_item(i):
+        for j in rules.get(i, ()):
+            if j in contains and j not in seen:
+                _fix_item(j)
+        fixed_update.append(i)
+        seen.add(i)
+
+    for i in update:
+        if i not in seen:
+            _fix_item(i)
+
+    return tuple(fixed_update)

solve.py

@@ -31,13 +31,15 @@ def main():
 
     try:
         print("Testing part 1 solution...")
+        test_input_p1 = getattr(mod, "test_input_p1", getattr(mod, "test_input", None))
         start = time()
-        assert mod.solve_p1(mod.test_input_p1) == mod.test_solution_p1
+        assert mod.solve_p1(test_input_p1) == mod.test_solution_p1
         print(f"Test passed in {time() - start:.3f} seconds")
 
         print("Testing part 2 solution...")
+        test_input_p2 = getattr(mod, "test_input_p2", getattr(mod, "test_input", None))
         start = time()
-        assert mod.solve_p2(mod.test_input_p2) == mod.test_solution_p2
+        assert mod.solve_p2(test_input_p2) == mod.test_solution_p2
         print(f"Test passed in {time() - start:.3f} seconds")
 
     except AssertionError: