Day 5: part 2

Author: horothesun

src/main/scala/Day05.scala

@@ -1,7 +1,8 @@
 import cats.Eq
-import cats.data.NonEmptyList
+import cats.data.{NonEmptyList, NonEmptyVector}
 import cats.derived.*
 import cats.syntax.all.*
+import scala.annotation.tailrec
 
 object Day05:
 
@@ -21,6 +22,10 @@ object Day05:
 
   case class Update(pages: NonEmptyList[Page]):
 
+    val toNev: NonEmptyVector[Page] = pages.toNev
+
+    val toNevWithIndex: NonEmptyVector[(Page, Int)] = toNev.zipWithIndex
+
     def allOrderRules: List[OrderRule] = pages.toList.zip(pages.tail).map(OrderRule.apply)
 
     def middlePage: Page =
@@ -30,6 +35,18 @@ object Day05:
     def firstViolatedOrderRules(rules: NonEmptyList[OrderRule]): Option[OrderRule] =
       allOrderRules.map(_.reverse).find(rules.contains_)
 
+    def swap(p1: Page, p2: Page): Update =
+      (toNevWithIndex.find(_._1 == p1), toNevWithIndex.find(_._1 == p2)).tupled.fold(ifEmpty = this) {
+        case ((_, i1), (_, i2)) => Update(pages = toNev.updatedUnsafe(i1, p2).updatedUnsafe(i2, p1).toNonEmptyList)
+      }
+
+    @tailrec
+    final def fixed(rules: NonEmptyList[OrderRule]): Update =
+      firstViolatedOrderRules(rules) match {
+        case Some(r) => swap(r.before, r.after).fixed(rules)
+        case None    => this
+      }
+
   object Update:
     def parse(s: String): Option[Update] = s.split(',').toList.toNel.flatMap(_.traverse(Page.parse).map(Update.apply))
 
@@ -47,6 +64,14 @@ object Day05:
         .collect(u => filteredRules(u).toNel.flatMap(u.firstViolatedOrderRules) match { case None => u })
         .foldMap(_.middlePage.n)
 
+    def fixedIncorrectlyOrderedUpdatesMiddlePageSum: Int =
+      updates.collect { u =>
+        filteredRules(u).toNel.flatMap(rs => u.firstViolatedOrderRules(rs).map((rs, _))) match {
+          case Some((rulesForUpdate, _)) => (u, rulesForUpdate)
+        }
+      }.map((u, rulesForUpdate) => u.fixed(rulesForUpdate))
+        .foldMap(_.middlePage.n)
+
   object Input:
     def parse(rows: List[String]): Option[Input] =
       rows.splitBySeparator("") match {

src/test/scala/Day05Suite.scala

@@ -118,6 +118,20 @@ class Day05Suite extends ScalaCheckSuite:
       5_064.some
     )
 
+  // part 2
+
+  test("small input incorrectly ordered updates middle page sum is 123"):
+    assertEquals(
+      Input.parse(smallInput).map(_.fixedIncorrectlyOrderedUpdatesMiddlePageSum),
+      123.some
+    )
+
+  test("big input incorrectly ordered updates middle page sum is 5_152"):
+    assertEquals(
+      Input.parse(bigInput).map(_.fixedIncorrectlyOrderedUpdatesMiddlePageSum),
+      5_152.some
+    )
+
 object Day05Suite:
 
   val bigInput: List[String] = getLinesFromFile("src/test/scala/day05_input.txt")