Add all paths Dijkstra solution to 2024 day 16 part 2

Author: sim642

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

@@ -34,36 +34,58 @@ object Day16 {
     Dijkstra.search(graphSearch).target.get._2
   }
 
-  def bestPathTiles(grid: Grid[Char]): Int = {
-    val forwardSearch = forwardGraphSearch(grid)
-    val forwardResult = Dijkstra.search(forwardSearch)
-
-    val backwardTraversal = new GraphTraversal[Reindeer] with UnitNeighbors[Reindeer] {
-      override val startNode: Reindeer = forwardResult.target.get._1 // TODO: other orientations
-
-      override def unitNeighbors(reindeer: Reindeer): IterableOnce[Reindeer] = {
-        val distance = forwardResult.distances(reindeer)
-        for {
-          (oldReindeer, step) <- Seq(
-            reindeer.copy(pos = reindeer.pos - reindeer.direction) -> 1, // backward steo
-            reindeer.copy(direction = reindeer.direction.left) -> 1000,
-            reindeer.copy(direction = reindeer.direction.right) -> 1000,
-          )
-          if grid(oldReindeer.pos) != '#'
-          oldDistance <- forwardResult.distances.get(oldReindeer)
-          if oldDistance + step == distance // if step on shortest path
-        } yield oldReindeer
+  trait Part2Solution {
+    def bestPathTiles(grid: Grid[Char]): Int
+  }
+
+  object BackwardNeighborsPart2Solution extends Part2Solution {
+    override def bestPathTiles(grid: Grid[Char]): Int = {
+      val forwardSearch = forwardGraphSearch(grid)
+      val forwardResult = Dijkstra.search(forwardSearch)
+
+      val backwardTraversal = new GraphTraversal[Reindeer] with UnitNeighbors[Reindeer] {
+        override val startNode: Reindeer = forwardResult.target.get._1 // TODO: other orientations
+
+        override def unitNeighbors(reindeer: Reindeer): IterableOnce[Reindeer] = {
+          val distance = forwardResult.distances(reindeer)
+          for {
+            (oldReindeer, step) <- Seq(
+              reindeer.copy(pos = reindeer.pos - reindeer.direction) -> 1, // backward step
+              reindeer.copy(direction = reindeer.direction.left) -> 1000,
+              reindeer.copy(direction = reindeer.direction.right) -> 1000,
+            )
+            if grid(oldReindeer.pos) != '#'
+            oldDistance <- forwardResult.distances.get(oldReindeer)
+            if oldDistance + step == distance // if step on shortest path
+          } yield oldReindeer
+        }
       }
+
+      BFS.traverse(backwardTraversal).nodes.map(_.pos).size
     }
+  }
 
-    BFS.traverse(backwardTraversal).nodes.map(_.pos).size
+  object AllPathsPart2Solution extends Part2Solution {
+    override def bestPathTiles(grid: Grid[Char]): Int = {
+      val forwardSearch = forwardGraphSearch(grid)
+      val forwardResult = Dijkstra.searchAllPaths(forwardSearch)
+
+      val backwardTraversal = new GraphTraversal[Reindeer] with UnitNeighbors[Reindeer] {
+        override val startNode: Reindeer = forwardResult.target.get._1 // TODO: other orientations
+
+        override def unitNeighbors(reindeer: Reindeer): IterableOnce[Reindeer] = forwardResult.allPrevNodes.getOrElse(reindeer, Set.empty)
+      }
+
+      BFS.traverse(backwardTraversal).nodes.map(_.pos).size
+    }
   }
 
   def parseGrid(input: String): Grid[Char] = input.linesIterator.map(_.toVector).toVector
 
   lazy val input: String = scala.io.Source.fromInputStream(getClass.getResourceAsStream("day16.txt")).mkString.trim
 
   def main(args: Array[String]): Unit = {
+    import AllPathsPart2Solution._
     println(lowestScore(parseGrid(input)))
     println(bestPathTiles(parseGrid(input)))
   }

src/main/scala/eu/sim642/adventofcodelib/graph/Dijkstra.scala

@@ -79,4 +79,58 @@ object Dijkstra {
       override def target: Option[(A, Int)] = None
     }
   }
+
+  // copied from search, modified like BFS.searchPaths
+  def searchAllPaths[A](graphSearch: GraphSearch[A]): Distances[A] & AllPaths[A] & Target[A] = {
+    val visitedDistance: mutable.Map[A, Int] = mutable.Map.empty
+    val prevNode: mutable.Map[A, mutable.Set[A]] = mutable.Map.empty
+    val toVisit: mutable.PriorityQueue[(Int, Option[A], A)] = mutable.PriorityQueue.empty(Ordering.by(-_._1))
+
+    def enqueue(oldNode: Option[A], node: A, dist: Int): Unit = {
+      toVisit.enqueue((dist, oldNode, node))
+    }
+
+    enqueue(None, graphSearch.startNode, 0)
+
+    while (toVisit.nonEmpty) {
+      val (dist, oldNode, node) = toVisit.dequeue()
+      if (!visitedDistance.contains(node)) {
+        visitedDistance(node) = dist
+        for (oldNode <- oldNode)
+          prevNode(node) = mutable.Set(oldNode)
+
+        if (graphSearch.isTargetNode(node, dist)) {
+          return new Distances[A] with AllPaths[A] with Target[A] {
+            override def distances: collection.Map[A, Int] = visitedDistance
+
+            override def allPrevNodes: collection.Map[A, collection.Set[A]] = prevNode
+
+            override def target: Option[(A, Int)] = Some(node -> dist)
+          }
+        }
+
+
+        def goNeighbor(newNode: A, distDelta: Int): Unit = {
+          if (!visitedDistance.contains(newNode)) { // avoids some unnecessary queue duplication but not all
+            val newDist = dist + distDelta
+            enqueue(Some(node), newNode, newDist)
+          }
+        }
+
+        graphSearch.neighbors(node).iterator.foreach(goNeighbor.tupled)
+      }
+      else { // visitedDistance.contains(node)
+        for (oldNode <- oldNode if visitedDistance(node) == dist)
+          prevNode(node) += oldNode
+      }
+    }
+
+    new Distances[A] with AllPaths[A] with Target[A] {
+      override def distances: collection.Map[A, Int] = visitedDistance
+
+      override def allPrevNodes: collection.Map[A, collection.Set[A]] = prevNode
+
+      override def target: Option[(A, Int)] = None
+    }
+  }
 }

src/main/scala/eu/sim642/adventofcodelib/graph/GraphTraversal.scala

@@ -29,6 +29,10 @@ trait Paths[A] {
     )
 }
 
+trait AllPaths[A] { // does not extend Paths, because prevNodes is Map, not function
+  def allPrevNodes: collection.Map[A, collection.Set[A]]
+}
+
 trait Order[A] {
   def nodeOrder: collection.Seq[A]
 }

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

@@ -1,9 +1,17 @@
 package eu.sim642.adventofcode2024
 
-import Day16._
+import Day16.*
+import Day16Test.*
+import org.scalatest.Suites
 import org.scalatest.funsuite.AnyFunSuite
 
-class Day16Test extends AnyFunSuite {
+class Day16Test extends Suites(
+  new Part1Test,
+  new BackwardNeighborsPart2SolutionTest,
+  new AllPathsPart2SolutionTest,
+)
+
+object Day16Test {
 
   val exampleInput =
     """###############
@@ -41,21 +49,31 @@ class Day16Test extends AnyFunSuite {
       |#S#.............#
       |#################""".stripMargin
 
-  test("Part 1 examples") {
-    assert(lowestScore(parseGrid(exampleInput)) == 7036)
-    assert(lowestScore(parseGrid(exampleInput2)) == 11048)
-  }
+  class Part1Test extends AnyFunSuite {
+    test("Part 1 examples") {
+      assert(lowestScore(parseGrid(exampleInput)) == 7036)
+      assert(lowestScore(parseGrid(exampleInput2)) == 11048)
+    }
 
-  test("Part 1 input answer") {
-    assert(lowestScore(parseGrid(input)) == 73404)
+    test("Part 1 input answer") {
+      assert(lowestScore(parseGrid(input)) == 73404)
+    }
   }
 
-  test("Part 2 examples") {
-    assert(bestPathTiles(parseGrid(exampleInput)) == 45)
-    assert(bestPathTiles(parseGrid(exampleInput2)) == 64)
-  }
+  class Part2SolutionTest(part2Solution: Part2Solution) extends AnyFunSuite {
+    import part2Solution._
 
-  test("Part 2 input answer") {
-    assert(bestPathTiles(parseGrid(input)) == 449)
+    test("Part 2 examples") {
+      assert(bestPathTiles(parseGrid(exampleInput)) == 45)
+      assert(bestPathTiles(parseGrid(exampleInput2)) == 64)
+    }
+
+    test("Part 2 input answer") {
+      assert(bestPathTiles(parseGrid(input)) == 449)
+    }
   }
+
+  class BackwardNeighborsPart2SolutionTest extends Part2SolutionTest(BackwardNeighborsPart2Solution)
+
+  class AllPathsPart2SolutionTest extends Part2SolutionTest(AllPathsPart2Solution)
 }