Solve 2024 day 21 part 2

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

@@ -1,9 +1,12 @@
 package eu.sim642.adventofcode2024
 
 import eu.sim642.adventofcodelib.Grid
-import eu.sim642.adventofcodelib.graph.{BFS, GraphSearch, Heuristic, UnitNeighbors}
+import eu.sim642.adventofcodelib.graph.{BFS, GraphSearch, GraphTraversal, Heuristic, SimultaneousBFS, TargetNode, UnitNeighbors}
 import eu.sim642.adventofcodelib.pos.Pos
 import eu.sim642.adventofcodelib.GridImplicits.*
+import eu.sim642.adventofcodelib.box.Box
+
+import scala.collection.mutable
 
 object Day21 {
 
@@ -78,18 +81,96 @@ object Day21 {
     BFS.search(graphSearch).target.get._2
   }
 
-  def codeComplexity(code: Code): Int = {
+
+  // copied & modified from 2024 day 10
+  // TODO: extract to library?
+  def pathSearch[A](graphSearch: GraphSearch[A] & UnitNeighbors[A]): GraphSearch[List[A]] & UnitNeighbors[List[A]] = {
+    new GraphSearch[List[A]] with UnitNeighbors[List[A]] {
+      override val startNode: List[A] = List(graphSearch.startNode)
+
+      override def unitNeighbors(node: List[A]): IterableOnce[List[A]] =
+        graphSearch.unitNeighbors(node.head).iterator.map(_ :: node)
+
+      override def isTargetNode(node: List[A], dist: Int): Boolean = graphSearch.isTargetNode(node.head, dist)
+    }
+  }
+
+  private def keypadPaths(keypad: Grid[Char]): Map[(Char, Char), Set[Code]] = {
+    val box = Box(Pos.zero, Pos(keypad(0).size - 1, keypad.size - 1))
+    (for {
+      startPos <- box.iterator
+      if keypad(startPos) != ' '
+      targetPos <- box.iterator
+      if keypad(targetPos) != ' '
+    } yield {
+      val graphSearch = new GraphSearch[Pos] with UnitNeighbors[Pos] with TargetNode[Pos] {
+        override val startNode: Pos = startPos
+
+        override def unitNeighbors(pos: Pos): IterableOnce[Pos] =
+          Pos.axisOffsets.map(pos + _).filter(keypad.containsPos).filter(keypad(_) != ' ')
+
+        override val targetNode: Pos = targetPos
+      }
+      (keypad(targetPos), keypad(startPos)) -> // flipped because paths are reversed
+        SimultaneousBFS.search(pathSearch(graphSearch))
+          .nodes
+          .filter(_.head == targetPos)
+          .map(poss =>
+            (poss lazyZip poss.tail)
+              .map({ case (p2, p1) => directionalOffsets.find(_._2 == p1 - p2).get._1 })
+              .mkString
+          )
+          .toSet
+    }).toMap
+  }
+
+  private val numericPaths: Map[(Char, Char), Set[Code]] = keypadPaths(numericKeypad)
+  private val directionalPaths: Map[(Char, Char), Set[Code]] = keypadPaths(directionalKeypad)
+
+  //println(numericPaths)
+
+  def shortestSequenceLength2(code: Code, directionalKeypads: Int, i: Int = 0): Long = {
+
+    val memo = mutable.Map.empty[(Code, Int), Long]
+
+    def helper(code: Code, i: Int): Long = {
+      memo.getOrElseUpdate((code, i), {
+        //assert(directionalKeypads == 0)
+        code.foldLeft(('A', 0L))({ case ((prev, length), cur) =>
+          val newLength =
+            (for {
+              path <- if (i == 0) numericPaths((prev, cur)) else directionalPaths((prev, cur))
+              path2 = path + 'A'
+              len =
+                if (i == directionalKeypads)
+                  path2.length.toLong
+                else
+                  helper(path2, i + 1)
+            } yield len).min
+          (cur, length + newLength)
+        })._2
+      })
+    }
+
+    helper(code, 0)
+  }
+
+
+  def codeComplexity(code: Code, directionalKeypads: Int): Long = {
     val numericPart = code.dropRight(1).toInt
-    shortestSequenceLength(code) * numericPart
+    shortestSequenceLength2(code, directionalKeypads) * numericPart
   }
 
-  def sumCodeComplexity(codes: Seq[Code]): Int = codes.map(codeComplexity).sum
+  def sumCodeComplexity(codes: Seq[Code], directionalKeypads: Int): Long = codes.map(codeComplexity(_, directionalKeypads)).sum
 
   def parseCodes(input: String): Seq[Code] = input.linesIterator.toSeq
 
   lazy val input: String = scala.io.Source.fromInputStream(getClass.getResourceAsStream("day21.txt")).mkString.trim
 
   def main(args: Array[String]): Unit = {
-    println(sumCodeComplexity(parseCodes(input)))
+    println(sumCodeComplexity(parseCodes(input), 2))
+    println(sumCodeComplexity(parseCodes(input), 25))
+
+    // part 2: 1301407762 - too low (Int overflowed in shortestSequenceLength2)
   }
 }

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

@@ -13,10 +13,18 @@ class Day21Test extends AnyFunSuite {
       |379A""".stripMargin
 
   test("Part 1 examples") {
-    assert(sumCodeComplexity(parseCodes(exampleInput)) == 126384)
+    assert(shortestSequenceLength2("029A", 0) == "<A^A>^^AvvvA".length)
+    assert(shortestSequenceLength2("029A", 1) == "v<<A>>^A<A>AvA<^AA>A<vAAA>^A".length)
+    assert(shortestSequenceLength2("029A", 2) == "<vA<AA>>^AvAA<^A>A<v<A>>^AvA^A<vA>^A<v<A>^A>AAvA^A<v<A>A>^AAAvA<^A>A".length)
+
+    assert(sumCodeComplexity(parseCodes(exampleInput), 2) == 126384)
+  }
+
+  test("Part 1 input answer") {
+    assert(sumCodeComplexity(parseCodes(input), 2) == 157892)
   }
 
-  test("Part 1 input") {
-    assert(sumCodeComplexity(parseCodes(input)) == 157892)
+  test("Part 2 input answer") {
+    assert(sumCodeComplexity(parseCodes(input), 25) == 197015606336332L)
   }
 }