Solve 2024 day 5 part 2

Author: sim642

src/main/scala/eu/sim642/adventofcode2024/Day5.scala

@@ -1,5 +1,7 @@
 package eu.sim642.adventofcode2024
 
+import scala.annotation.tailrec
+
 object Day5 {
 
   type Rule = (Int, Int)
@@ -23,6 +25,33 @@ object Day5 {
       .sum
   }
 
+  def sort(rules: Seq[Rule], update: Update): Update = {
+    val updateRules = rules.filter(rule => update.contains(rule._1) && update.contains(rule._2))
+    val updateRulesMap = updateRules.toSet.groupMap(_._1)(_._2).withDefaultValue(Set.empty)
+
+    // TODO: optimize?
+    @tailrec
+    def topologicalSort(pages: Set[Int], updateRulesMap: Map[Int, Set[Int]], sortedUpdate: List[Int]): List[Int] = {
+      pages.find(updateRulesMap(_).isEmpty) match {
+        case None => sortedUpdate
+        case Some(before) =>
+          val newUpdateRulesMap = (updateRulesMap - before).view.mapValues(_ - before).toMap
+          topologicalSort(pages - before, newUpdateRulesMap, before :: sortedUpdate)
+      }
+    }
+
+    topologicalSort(update.toSet, updateRulesMap, Nil)
+  }
+
+  def sumIncorrectMiddles(input: Input): Int = {
+    val Input(rules, updates) = input
+    updates
+      .filter(!isCorrect(rules, _))
+      .map(sort(rules, _))
+      .map(update => update(update.size / 2))
+      .sum
+  }
+
   def parseRule(s: String): Rule = s match {
     case s"$x|$y" => (x.toInt, y.toInt)
   }
@@ -40,5 +69,6 @@ object Day5 {
 
   def main(args: Array[String]): Unit = {
     println(sumCorrectMiddles(parseInput(input)))
+    println(sumIncorrectMiddles(parseInput(input)))
   }
 }

src/test/scala/eu/sim642/adventofcode2024/Day5Test.scala

@@ -51,4 +51,18 @@ class Day5Test extends AnyFunSuite {
   test("Part 1 input answer") {
     assert(sumCorrectMiddles(parseInput(input)) == 6051)
   }
+
+  test("Part 2 examples") {
+    val input@Input(rules, updates) = parseInput(exampleInput)
+
+    assert(sort(rules, updates(3)) == parseUpdate("97,75,47,61,53"))
+    assert(sort(rules, updates(4)) == parseUpdate("61,29,13"))
+    assert(sort(rules, updates(5)) == parseUpdate("97,75,47,29,13"))
+
+    assert(sumIncorrectMiddles(input) == 123)
+  }
+
+  test("Part 2 input answer") {
+    assert(sumIncorrectMiddles(parseInput(input)) == 5093)
+  }
 }