BenScape BenScape - 26 days ago 11
Scala Question

Tail-recursive Implementation of Scala's List partition method

For exercise purposes I've been trying to implement a couple of Scala's List methods in a functional manner, one of them being

partition
. Assume the following signature:

def partition[T](l: List[T], f: T => Boolean): (List[T], List[T])


It returns a tuple consisting of two lists - the first one contains all the elements from
l
that fulfill the passed predicate
f
and another one which contains all the other elements.

I came up with the following recursive solution which is unfortunately not tail-recursive:

def partition[T](l: List[T], f: T => Boolean): (List[T], List[T]) = {
l match {
case Nil => (Nil, Nil)
case head :: rest => {
val (left, right) = partition(rest, f)
if (f(head))
(head :: left, right)
else
(left, head :: right)
}
}
}


In this stack overflow question (Can all recursive functions be re-written as tail-recursions?) it is made clear that an accumulator could be used in some cases. In the given one I would claim that this is not possible since it would return the final lists in a reversed manner.

Could you please give me a tail-recursive solution? Maybe even with continuation passing (I haven't really understood how it works and how it could be applied)?

Answer Source

You can keep a tuple as the accumulator, and make sure to reverse the lists before returning them:

def partition[T](l: List[T])(f: T => Boolean): (List[T], List[T]) = {
  @tailrec
  def partitionInternal(l: List[T])(acc: (List[T], List[T])): (List[T], List[T]) = {
    l match {
      case Nil => acc
      case head :: tail =>
        if (f(head)) partitionInternal(tail)(head :: acc._1, acc._2)
        else partitionInternal(tail)(acc._1, head :: acc._2)
    }
  }

  val (lf, r) = partitionInternal(l)((List.empty[T], List.empty[T]))
  (lf.reverse, r.reverse)
}

An alternative solution would be to append (:+) instead of prepending (::), but then you pay the price of O(n) for every entry.

Another idea would be to use a ListBuffer[T] instead of List[T] for the internal recursive implementation which has constant time append. All you need to do is call .toList at the end:

def partition[T](l: List[T])(f: T => Boolean): (List[T], List[T]) = {
  @tailrec
  def partitionInternal(l: List[T])(acc: (ListBuffer[T], ListBuffer[T])): (ListBuffer[T], ListBuffer[T])  = {
    l match {
      case Nil => acc
      case head :: tail =>
        val (leftAcc, rightAcc) = acc
        if (f(head)) partitionInternal(tail)((leftAcc += head, rightAcc))
        else partitionInternal(tail)((leftAcc, rightAcc += head))
    }
  }

  val (lf, r) = partitionInternal(l)((ListBuffer.empty[T], ListBuffer.empty[T]))
  (lf.toList, r.toList)
}

Additionaly, notice that I create a separate argument list for the List[T] and the function from T => Boolean. That is in order to help the compiler infer the right type argument when applying the method since type inference flows between parameter lists.