Tail-recursive Implementation of Scala's List partition method

Question

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)?

You can reverse the results before returning them. – Dima Sep 26 '17 at 15:23 — Dima, Sep 26 '17 at 15:23

score 1 · Answer 1 · answered Sep 26 '17 at 15:27

You need to keep two accumulators, one for left and one for right. When you're done going through the input list, simply return both accumulators (reversing them to get back to the original order):

def partition[T](l: List[T], f: T => Boolean): (List[T], List[T]) = {
  @annotation.tailrec
  def aux(tl: List[T], left: List[T], right: List[T]): (List[T], List[T]) = tl match {
    case Nil => (left.reverse, right.reverse)
    case head :: rest => {
      if (f(head))
        aux(rest, head :: left, right)
      else
        aux(rest, left, head :: right)
    }
  }

  aux(l, List(), List())
}

Using it:

scala> def partition[T](l: List[T], f: T => Boolean): (List[T], List[T]) = {
     |   @annotation.tailrec
     |   def aux(tl: List[T], left: List[T], right: List[T]): (List[T], List[T]) = tl match {
     |     case Nil => (left.reverse, right.reverse)
     |     case head :: rest => {
     |       if (f(head))
     |         aux(rest, head :: left, right)
     |       else
     |         aux(rest, left, head :: right)
     |     }
     |   }
     | 
     |   aux(l, List(), List())
     | } 
partition: [T](l: List[T], f: T => Boolean)(List[T], List[T])

scala> partition(List(1, 2, 3, 4, 5), (i: Int) => i%2 == 0)
res1: (List[Int], List[Int]) = (List(2, 4),List(1, 3, 5))

Yuval Itzchakov · Accepted Answer · 2017-09-26T15:45:31.187

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.

Tail-recursive Implementation of Scala's List partition method

2 Answers2