From bbe3ba8644519db7040da93d3ff7c3dc3c268040 Mon Sep 17 00:00:00 2001
From: Bert Peters <bert@bertptrs.nl>
Date: Thu, 5 Dec 2024 08:46:40 +0100
Subject: [PATCH] Implement 2024 day 5

---
 2024/src/aoc/days/day5.py | 50 +++++++++++++++++++++++++++++++++++++++
 2024/tests/samples/05.txt | 28 ++++++++++++++++++++++
 2024/tests/test_day5.py   | 19 +++++++++++++++
 3 files changed, 97 insertions(+)
 create mode 100644 2024/src/aoc/days/day5.py
 create mode 100644 2024/tests/samples/05.txt
 create mode 100644 2024/tests/test_day5.py

diff --git a/2024/src/aoc/days/day5.py b/2024/src/aoc/days/day5.py
new file mode 100644
index 0000000..cd9c30f
--- /dev/null
+++ b/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
diff --git a/2024/tests/samples/05.txt b/2024/tests/samples/05.txt
new file mode 100644
index 0000000..9d146d6
--- /dev/null
+++ b/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
diff --git a/2024/tests/test_day5.py b/2024/tests/test_day5.py
new file mode 100644
index 0000000..b00e7db
--- /dev/null
+++ b/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