diff --git a/src/main/scala/eu/sim642/adventofcode2024/Day18.scala b/src/main/scala/eu/sim642/adventofcode2024/Day18.scala
index 1294ea7b..b3126f53 100644
--- a/src/main/scala/eu/sim642/adventofcode2024/Day18.scala
+++ b/src/main/scala/eu/sim642/adventofcode2024/Day18.scala
@@ -57,7 +57,7 @@ object Day18 {
   object LinearOnPathPart2Solution extends Part2Solution {
     def exitPath(bytes: Seq[Pos], max: Pos, after: Int): Option[Seq[Pos]] = {
       val graphSearch = bytesGraphSearch(bytes, max, after + 1)
-      BFS.searchPaths(graphSearch).paths.get(graphSearch.targetNode) // TODO: optimize paths to not compute everything
+      BFS.searchPaths(graphSearch).paths.lift(graphSearch.targetNode)
     }
 
     override def findBlockingByte(bytes: Seq[Pos], max: Pos = Pos(70, 70)): Pos = {
diff --git a/src/main/scala/eu/sim642/adventofcodelib/graph/GraphTraversal.scala b/src/main/scala/eu/sim642/adventofcodelib/graph/GraphTraversal.scala
index 6d1559cd..6798d97f 100644
--- a/src/main/scala/eu/sim642/adventofcodelib/graph/GraphTraversal.scala
+++ b/src/main/scala/eu/sim642/adventofcodelib/graph/GraphTraversal.scala
@@ -23,14 +23,10 @@ trait Distances[A] {
 trait Paths[A] {
   def prevNodes: collection.Map[A, A]
 
-  def paths: collection.Map[A, Seq[A]] = {
-    prevNodes.map((node, _) =>
-      node -> (node +: LazyList.unfold0(node)(prevNodes.get)).reverse
+  def paths: PartialFunction[A, Seq[A]] =
+    prevNodes.andThen(node =>
+      (node #:: LazyList.unfold0(node)(prevNodes.get)).reverse // TODO: don't bother with LazyList now that it's a function
     )
-  }
-
-  /*def paths(node: A): Seq[A] =
-    (node +: LazyList.unfold0(node)(prevNodes.get)).reverse*/
 }
 
 trait Order[A] {