From f3f5fbcbdd0ddef43b6111f4e501ed2a1d6bd169 Mon Sep 17 00:00:00 2001
From: Jeremy Kaplan <code@jdkaplan.dev>
Date: Wed, 4 Dec 2024 22:27:10 -0800
Subject: [PATCH] 2024 Day 5 Part 2

---
 2024/src/day05.py | 58 +++++++++++++++++++++++++++++++----------------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/2024/src/day05.py b/2024/src/day05.py
index fdf58f7..50db683 100644
--- a/2024/src/day05.py
+++ b/2024/src/day05.py
@@ -1,6 +1,7 @@
 import os
 from collections import defaultdict
 from dataclasses import dataclass
+from functools import cmp_to_key
 
 
 def relpath(path: str) -> str:
@@ -28,26 +29,43 @@ class Constraint:
     after: set[int]
 
 
-def part1(rules, updates):
-    constraints = defaultdict(lambda: Constraint(before=set(), after=set()))
+constraints = defaultdict(lambda: Constraint(before=set(), after=set()))
 
-    for (x, y) in rules:
-        constraints[x].after.add(y)
-        constraints[y].before.add(x)
-
-    middles = 0
-    for update in updates:
-        valid = True
-        for (x, y) in zip(update[:-1], update[1:]):
-            if x not in constraints or y in constraints[x].after:
-                continue
-            valid = False
-            break
-
-        if valid:
-            middles += update[int(len(update)) // 2]
-
-    return middles
+for (x, y) in rules:
+    constraints[x].after.add(y)
+    constraints[y].before.add(x)
 
 
-print(part1(rules, updates))
+def is_valid(update):
+    for (x, y) in zip(update[:-1], update[1:]):
+        if x not in constraints or y in constraints[x].after:
+            continue
+        return False
+    return True
+
+
+def part1():
+    return sum(
+        update[int(len(update)) // 2]
+        for update in updates
+        if is_valid(update)
+    )
+
+
+def part2():
+    def compare(x, y):
+        if y in constraints[x].after:
+            return +1
+        if y in constraints[x].before:
+            return -1
+        return 0
+
+    return sum(
+        sorted(update, key=cmp_to_key(compare))[int(len(update)) // 2]
+        for update in updates
+        if not is_valid(update)
+    )
+
+
+print(part1())
+print(part2())