Extract and test solutions in 2024 day 21

sim642 · sim642 · commit 0fe17f7eec95 · 2024-12-21T12:18:00.000+02:00
diff --git a/src/main/scala/eu/sim642/adventofcode2024/Day21.scala b/src/main/scala/eu/sim642/adventofcode2024/Day21.scala
@@ -31,145 +31,154 @@ object Day21 {
     '<' -> Pos(-1, 0),
   )
 
-  case class State(directionalPoss: List[Pos], numericPos: Pos, input: Code) {
-
-    def numericPress(button: Char): Option[State] = button match {
-      case 'A' =>
-        val newButton = numericKeypad(numericPos)
-        Some(copy(input = input + newButton))
-      case _ =>
-        val offset = directionalOffsets(button)
-        val newNumericPos = numericPos + offset
-        if (numericKeypad.containsPos(newNumericPos) && numericKeypad(newNumericPos) != ' ')
-          Some(copy(numericPos = newNumericPos))
-        else
-          None // out of keypad
-    }
+  trait Solution {
+    def shortestSequenceLength(code: Code, directionalKeypads: Int): Long
 
-    def directionalPress(button: Char): Option[State] = directionalPoss match {
-      case Nil => numericPress(button)
-      case directionalPos :: newDirectionalPoss =>
-        button match {
-          case 'A' =>
-            val newButton = directionalKeypad(directionalPos)
-            copy(directionalPoss = newDirectionalPoss).directionalPress(newButton).map(newState =>
-              newState.copy(directionalPoss = directionalPos :: newState.directionalPoss)
-            )
-          case _ =>
-            val offset = directionalOffsets(button)
-            val newDirectionalPos = directionalPos + offset
-            if (directionalKeypad.containsPos(newDirectionalPos) && directionalKeypad(newDirectionalPos) != ' ')
-              Some(copy(directionalPoss = newDirectionalPos :: newDirectionalPoss))
-            else
-              None // out of keypad
-        }
+    def codeComplexity(code: Code, directionalKeypads: Int): Long = {
+      val numericPart = code.dropRight(1).toInt
+      shortestSequenceLength(code, directionalKeypads) * numericPart
     }
 
-    def userPress(button: Char): Option[State] = directionalPress(button)
+    def sumCodeComplexity(codes: Seq[Code], directionalKeypads: Int): Long = codes.map(codeComplexity(_, directionalKeypads)).sum
   }
 
-  def shortestSequenceLength(code: Code): Int = {
-
-    val graphSearch = new GraphSearch[State] with UnitNeighbors[State] {
-      override val startNode: State = State(List.fill(2)(directionalKeypad.posOf('A')), numericKeypad.posOf('A'), "")
+  object NaiveSolution extends Solution {
+
+    case class State(directionalPoss: List[Pos], numericPos: Pos, input: Code) {
+
+      def numericPress(button: Char): Option[State] = button match {
+        case 'A' =>
+          val newButton = numericKeypad(numericPos)
+          Some(copy(input = input + newButton))
+        case _ =>
+          val offset = directionalOffsets(button)
+          val newNumericPos = numericPos + offset
+          if (numericKeypad.containsPos(newNumericPos) && numericKeypad(newNumericPos) != ' ')
+            Some(copy(numericPos = newNumericPos))
+          else
+            None // out of keypad
+      }
 
-      override def unitNeighbors(state: State): IterableOnce[State] = "<v>^A".iterator.flatten(state.userPress).filter(s => code.startsWith(s.input))
+      def directionalPress(button: Char): Option[State] = directionalPoss match {
+        case Nil => numericPress(button)
+        case directionalPos :: newDirectionalPoss =>
+          button match {
+            case 'A' =>
+              val newButton = directionalKeypad(directionalPos)
+              copy(directionalPoss = newDirectionalPoss).directionalPress(newButton).map(newState =>
+                newState.copy(directionalPoss = directionalPos :: newState.directionalPoss)
+              )
+            case _ =>
+              val offset = directionalOffsets(button)
+              val newDirectionalPos = directionalPos + offset
+              if (directionalKeypad.containsPos(newDirectionalPos) && directionalKeypad(newDirectionalPos) != ' ')
+                Some(copy(directionalPoss = newDirectionalPos :: newDirectionalPoss))
+              else
+                None // out of keypad
+          }
+      }
 
-      override def isTargetNode(state: State, dist: Int): Boolean = state.input == code
+      def userPress(button: Char): Option[State] = directionalPress(button)
     }
 
-    BFS.search(graphSearch).target.get._2
-  }
+    override def shortestSequenceLength(code: Code, directionalKeypads: Int): Long = {
 
+      val graphSearch = new GraphSearch[State] with UnitNeighbors[State] {
+        override val startNode: State = State(List.fill(directionalKeypads)(directionalKeypad.posOf('A')), numericKeypad.posOf('A'), "")
 
-  // 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(state: State): IterableOnce[State] = "<v>^A".iterator.flatten(state.userPress).filter(s => code.startsWith(s.input))
 
-      override def unitNeighbors(node: List[A]): IterableOnce[List[A]] =
-        graphSearch.unitNeighbors(node.head).iterator.map(_ :: node)
+        override def isTargetNode(state: State, dist: Int): Boolean = state.input == code
+      }
 
-      override def isTargetNode(node: List[A], dist: Int): Boolean = graphSearch.isTargetNode(node.head, dist)
+      BFS.search(graphSearch).target.get._2
     }
   }
 
-  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
-  }
+  object DynamicProgrammingSolution extends Solution {
+
+    // 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)
 
-  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
-      })
+        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)
+      }
     }
 
-    helper(code, 0)
-  }
+    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)
+
+    override def shortestSequenceLength(code: Code, directionalKeypads: Int): 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
+        })
+      }
 
-  def codeComplexity(code: Code, directionalKeypads: Int): Long = {
-    val numericPart = code.dropRight(1).toInt
-    shortestSequenceLength2(code, directionalKeypads) * numericPart
+      helper(code, 0)
+    }
   }
 
-  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
 
+  val part1DirectionalKeypads = 2
+  val part2DirectionalKeypads = 25
+
   def main(args: Array[String]): Unit = {
-    println(sumCodeComplexity(parseCodes(input), 2))
-    println(sumCodeComplexity(parseCodes(input), 25))
+    import DynamicProgrammingSolution._
+    println(sumCodeComplexity(parseCodes(input), part1DirectionalKeypads))
+    println(sumCodeComplexity(parseCodes(input), part2DirectionalKeypads))
 
     // part 2: 1301407762 - too low (Int overflowed in shortestSequenceLength2)
   }
diff --git a/src/test/scala/eu/sim642/adventofcode2024/Day21Test.scala b/src/test/scala/eu/sim642/adventofcode2024/Day21Test.scala
@@ -1,9 +1,16 @@
 package eu.sim642.adventofcode2024
 
-import Day21._
+import Day21.*
+import Day21Test.*
+import org.scalatest.Suites
 import org.scalatest.funsuite.AnyFunSuite
 
-class Day21Test extends AnyFunSuite {
+class Day21Test extends Suites(
+  new NaiveSolutionTest,
+  new DynamicProgrammingSolutionTest,
+)
+
+object Day21Test {
 
   val exampleInput =
     """029A
@@ -12,19 +19,37 @@ class Day21Test extends AnyFunSuite {
       |456A
       |379A""".stripMargin
 
-  test("Part 1 examples") {
-    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)
+  abstract class SolutionTest(solution: Solution) extends AnyFunSuite {
+    import solution._
 
-    assert(sumCodeComplexity(parseCodes(exampleInput), 2) == 126384)
-  }
+    test("Part 1 examples") {
+      assert(shortestSequenceLength("029A", 0) == "<A^A>^^AvvvA".length)
+      assert(shortestSequenceLength("029A", 1) == "v<<A>>^A<A>AvA<^AA>A<vAAA>^A".length)
+      assert(shortestSequenceLength("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), part1DirectionalKeypads) == 126384)
+    }
+
+    test("Part 1 input answer") {
+      assert(sumCodeComplexity(parseCodes(input), part1DirectionalKeypads) == 157892)
+    }
 
-  test("Part 1 input answer") {
-    assert(sumCodeComplexity(parseCodes(input), 2) == 157892)
+    protected val testPart2: Boolean = true
+
+    if (testPart2) {
+      test("Part 2 examples") {
+        assert(sumCodeComplexity(parseCodes(exampleInput), part2DirectionalKeypads) == 154115708116294L) // not in text
+      }
+
+      test("Part 2 input answer") {
+        assert(sumCodeComplexity(parseCodes(input), part2DirectionalKeypads) == 197015606336332L)
+      }
+    }
   }
 
-  test("Part 2 input answer") {
-    assert(sumCodeComplexity(parseCodes(input), 25) == 197015606336332L)
+  class NaiveSolutionTest extends SolutionTest(NaiveSolution) {
+    override protected val testPart2: Boolean = false
   }
+
+  class DynamicProgrammingSolutionTest extends SolutionTest(DynamicProgrammingSolution)
 }
