Cartesian product in Scala

Question

I'm using this code to compute a cartesian product:

def cartesianProduct[T](in: Seq[Seq[T]]): Seq[Seq[T]] = {
  def loop(acc: Seq[T], rest: Seq[Seq[T]]): Seq[Seq[T]] = {
    rest match {
      case seq :: Nil => seq.map(acc :+ _)
      case seq :: remainingSeqs => seq.flatMap(i => loop(acc :+ i, remainingSeqs))
    }
  }

  loop(Seq(), in)
}

To give an example, it transorfoms this input into this output:

scala> cartesianProduct(Seq(
  Seq(1, 2, 3),
  Seq(4),
  Seq(5, 6)
))
res0: Seq[Seq[Int]] = List(
  List(1, 4, 5), List(1, 4, 6), 
  List(2, 4, 5), List(2, 4, 6), 
  List(3, 4, 5), List(3, 4, 6)
)

In this method, the loop function is not tail-recursive and cannot be since it is used from the flatMap.

Can this be re-written to be tail-recursive?

Answer 1

It can actually be re-written as a tail-recursive function. The key was to use a global accumulator (with all results) and not a "local" accumlator per dimension.

Here is a working tail-recursive version of this function:

def cartesianProduct[T](in: Seq[Seq[T]]): Seq[Seq[T]] = {
  @scala.annotation.tailrec
  def loop(acc: Seq[Seq[T]], rest: Seq[Seq[T]]): Seq[Seq[T]] = {
    rest match {
      case Nil => 
        acc
      case seq :: remainingSeqs => 
        // Equivalent of: 
        // val next = seq.flatMap(i => acc.map(a => i+: a))
        val next = for {
          i <- seq
          a <- acc
        } yield i +: a
        loop(next, remainingSeqs)
    }
  }

  loop(Seq(Nil), in.reverse)
}