day5 part2

2024/python/day05.py

@@ -14,30 +14,53 @@ def read(rows):
             rules.append(list(map(int, row.split("|"))))
     return rules, updates
 
-
-# @shared.profile
-def part1(rows):
-    rules, updates = read(rows)
-    bad = []
+def get_bads(rules, updates):
+    bads = []
     for idx, update in enumerate(updates):
         for first, second in rules:
             if first in update and second in update:    
                 first_idx = update.index(first)
                 second_idx = update.index(second)
                 if first_idx > second_idx:
-                    bad.append(idx)
+                    bads.append(idx)
                     break
+    return bads
+
+# @shared.profile
+def part1(rows):
+    rules, updates = read(rows)
+    bads = get_bads(rules, updates)
     total = 0
     for idx, update in enumerate(updates):
-        if idx in bad:
+        if idx in bads:
             continue
         total += update[len(update) // 2]
     print(total)
 
+def fix_bad(rules, bad):
+    for first, second in rules:
+        if first in bad and second in bad: 
+            first_idx = bad.index(first)
+            second_idx = bad.index(second)
+            if first_idx > second_idx:
+                bad[first_idx], bad[second_idx] = bad[second_idx], bad[first_idx]
+    return bad
 
 # @shared.profile
 def part2(rows):
-    pass
+    rules, updates = read(rows)
+    bads = get_bads(rules, updates)
+    total = 0
+    for idx in bads:
+        bad = updates[idx]
+        while True:
+            after = fix_bad(rules, bad[:])
+            if bad == after:
+                break
+            bad = after
+        total += bad[len(bad) // 2]
+    print(total)
+
 
 def main():
     rows = [row for row in shared.load_rows(5)]