Implement 2024 day 5

2024/src/aoc/days/day5.py

@@ -0,0 +1,50 @@
+from collections import defaultdict
+import functools
+
+from . import CombinedRunner
+
+
+def parse_input(input: str) -> tuple[set[tuple[int, int]], list[list[int]]]:
+    first, second = input.strip().split("\n\n")
+
+    rules = {tuple(int(x) for x in line.split("|")) for line in first.split("\n")}
+    updates = [[int(x) for x in line.split(",")] for line in second.split("\n")]
+
+    return rules, updates
+
+
+def is_correct(update: list[int], must_after: dict[int, set[int]]) -> bool:
+    forbidden = set()
+
+    for entry in update:
+        if entry in forbidden:
+            return False
+
+        forbidden |= must_after.get(entry, set())
+        
+    return True
+
+
+class DayRunner(CombinedRunner):
+    @classmethod
+    def run_both(cls, input: str) -> int:
+        rules, updates = parse_input(input)
+
+        must_after = defaultdict(set)
+
+        for before, after in rules:
+            must_after[after].add(before)
+
+        correct = 0
+        corrected = 0
+
+        key = functools.cmp_to_key(lambda a, b: -1 if (a, b) in rules else 1)
+
+        for update in updates:
+            if is_correct(update, must_after):
+                correct += update[len(update) // 2]
+            else:
+                update.sort(key=key)
+                corrected += update[len(update) // 2]
+
+        return correct, corrected

2024/tests/samples/05.txt

@@ -0,0 +1,28 @@
+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

2024/tests/test_day5.py

@@ -0,0 +1,19 @@
+import os
+
+from aoc.days.day5 import DayRunner
+
+def get_data() -> str:
+    sample = os.path.dirname(__file__) + "/samples/05.txt"
+    with open(sample, mode="rt", encoding="utf-8") as f:
+        return f.read()
+
+
+def test_sample_part1() -> None:
+    data = get_data()
+
+    assert DayRunner.part1(data) == 143
+
+def test_sample_part2() -> None:
+    data = get_data()
+
+    assert DayRunner.part2(data) == 123