zipWith function in Scala

Question

My implementation works, but it looks very ugly. Curious to see how it's better implemented.

The zipWith function takes a list of lists (List[List[A]]), or a row with columns if you like. It will then transform this into a column with rows and applies a function List[A]=>B to obtain a List[B]. Hopefully the unit test below explains this better.

  test("zipWith"){
    val input = List(List(1,2,3), List(4,5,6), List(7,8,9))
    val expected = List("1::4::7", "2::5::8", "3::6::9")
    assert(zipWith(input)(_.mkString("::")) === expected)
  }

and the implementation is:

  def zipWith[A, B](lists: List[List[A]])(f: List[A] => B): List[B] = {
    @tailrec
    def loop(acc: List[B], input: List[List[A]]): List[B] = {
      val init = (Nil: List[List[A]], Nil: List[A])
      val (left, zipped) = input.foldLeft(init)((z, a) => {
        val (tails, heads) = z
        (a.tail :: tails, a.head :: heads)
      })
      if (left.foldLeft(false)(_ || _.isEmpty))
        f(zipped.reverse) :: acc
      else
        loop(f(zipped.reverse) :: acc, left.reverse)
    }

    loop(Nil, lists).reverse
  }

Answer 1

How about this?

def zipWith[A, B](ts: Traversable[Traversable[A]])(f: Traversable[A] => B): Traversable[B] = {
  ts.transpose.map(f)
}

Example:

scala> zipWith(input)(_.mkString("::"))
res10: Traversable[String] = List(1::4::7, 2::5::8, 3::6::9)

Answer 2

Two things:

• Your variable names could be better, I would rather use rows and columns as a reader can understand it easier.
• List is an implementation of Seq, you should probably use Seq where you can.

With that said I would do something like this:

def zipWith[A, B](lists: Seq[Seq[A]])(f: Seq[A] => B): Seq[B] = {
  lists.headOption.map(_.size) match {
    case None =>
      // Input list is empty, return empty
      Seq.empty[B]
    case Some(columnCount) =>
      // Do a sanity check, based on column length of first row.
      require(lists.forall(_.size == columnCount), "Can't zip rows that have differing column lengths.")
      val initial = Seq.fill(columnCount)(Seq.empty[A])
      val zippedByColumns = lists.foldLeft(initial) {
        case (columns, nextRow) =>
          columns.zip(nextRow).map{ case (column, rowElement) => column :+ rowElement }
      }
      zippedByColumns.map(f)
  }
}