Scala foldLeft while some conditions are true

问题

How to emulate following behavior in Scala? i.e. keep folding while some certain conditions on the accumulator are met.

def foldLeftWhile[B](z: B, p: B => Boolean)(op: (B, A) => B): B

For example

scala> val seq = Seq(1, 2, 3, 4)
seq: Seq[Int] = List(1, 2, 3, 4)
scala> seq.foldLeftWhile(0, _ < 3) { (acc, e) => acc + e }
res0: Int = 1
scala> seq.foldLeftWhile(0, _ < 7) { (acc, e) => acc + e }
res1: Int = 6

UPDATES:

Based on @Dima answer, I realized that my intention was a little bit side-effectful. So I made it synchronized with takeWhile, i.e. there would be no advancement if the predicate does not match. And add some more examples to make it clearer. (Note: that will not work with Iterators)

回答1:

First, note that your example seems wrong. If I understand correctly what you describe, the result should be 1 (the last value on which the predicate _ < 3 was satisfied), not 6

The simplest way to do this is using a return statement, which is very frowned upon in scala, but I thought, I'd mention it for the sake of completeness.

def foldLeftWhile[A, B](seq: Seq[A], z: B, p: B => Boolean)(op: (B, A) => B): B = foldLeft(z) { case (b, a) => 
   val result = op(b, a) 
   if(!p(result)) return b
   result
}

Since we want to avoid using return, scanLeft might be a possibility:

seq.toStream.scanLeft(z)(op).takeWhile(p).last

This is a little wasteful, because it accumulates all (matching) results. You could use iterator instead of toStream to avoid that, but Iterator does not have .last for some reason, so, you'd have to scan through it an extra time explicitly:

 seq.iterator.scanLeft(z)(op).takeWhile(p).foldLeft(z) { case (_, b) => b }

回答2:

It is pretty straightforward to define what you want in scala. You can define an implicit class which will add your function to any TraversableOnce (that includes Seq).

implicit class FoldLeftWhile[A](trav: TraversableOnce[A]) {
  def foldLeftWhile[B](init: B)(where: B => Boolean)(op: (B, A) => B): B = {
    trav.foldLeft(init)((acc, next) => if (where(acc)) op(acc, next) else acc)
  }
}
Seq(1,2,3,4).foldLeftWhile(0)(_ < 3)((acc, e) => acc + e)

Update, since the question was modified:

implicit class FoldLeftWhile[A](trav: TraversableOnce[A]) {
  def foldLeftWhile[B](init: B)(where: B => Boolean)(op: (B, A) => B): B = {
    trav.foldLeft((init, false))((a,b) => if (a._2) a else {
      val r = op(a._1, b)
      if (where(r)) (op(a._1, b), false) else (a._1, true)
    })._1
  }
}

Note that I split your (z: B, p: B => Boolean) into two higher-order functions. That's just a personal scala style preference.

回答3:

What about this:

def foldLeftWhile[A, B](z: B, xs: Seq[A], p: B => Boolean)(op: (B, A) => B): B = {
  def go(acc: B, l: Seq[A]): B = l match {
    case h +: t => 
        val nacc = op(acc, h)
        if(p(nacc)) go(op(nacc, h), t) else nacc
    case _ => acc
  }
  go(z, xs)
}

val a = Seq(1,2,3,4,5,6)
val r = foldLeftWhile(0, a, (x: Int) => x <= 3)(_ + _)
println(s"$r")

Iterate recursively on the collection while the predicate is true, and then return the accumulator.

You cand try it on scalafiddle

回答4:

Fist exemple: predicate for each element

First you can use inner tail recursive function

implicit class TravExt[A](seq: TraversableOnce[A]) {
  def foldLeftWhile[B](z: B, f: A => Boolean)(op: (A, B) => B): B = {
    @tailrec
    def rec(trav: TraversableOnce[A], z: B): B = trav match {
      case head :: tail if f(head) => rec(tail, op(head, z))
      case _ => z
    }
    rec(seq, z)
  }
}

Or short version

implicit class TravExt[A](seq: TraversableOnce[A]) {
  @tailrec
  final def foldLeftWhile[B](z: B, f: A => Boolean)(op: (A, B) => B): B = seq match {
    case head :: tail if f(head) => tail.foldLeftWhile(op(head, z), f)(op)
    case _ => z
  }
}

Then use it

val a = List(1, 2, 3, 4, 5, 6).foldLeftWhile(0, _ < 3)(_ + _) //a == 3

Second example: for accumulator value:

implicit class TravExt[A](seq: TraversableOnce[A]) {
  def foldLeftWhile[B](z: B, f: A => Boolean)(op: (A, B) => B): B = {
    @tailrec
    def rec(trav: TraversableOnce[A], z: B): B = trav match {
      case _ if !f(z) => z
      case head :: tail => rec(tail, op(head, z))
      case _ => z
    }
    rec(seq, z)
  }
}

Or short version

implicit class TravExt[A](seq: TraversableOnce[A]) {
  @tailrec
  final def foldLeftWhile[B](z: B, f: A => Boolean)(op: (A, B) => B): B = seq match {
    case _ if !f(z) => z
    case head :: tail => tail.foldLeftWhile(op(head, z), f)(op)
    case _ => z
  }
}

回答5:

After a while I received a lot of good looking answers. So, I combined them to this single post

a very concise solution by @Dima

implicit class FoldLeftWhile[A](seq: Seq[A]) {

  def foldLeftWhile[B](z: B)(p: B => Boolean)(op: (B, A) => B): B = {
    seq.toStream.scanLeft(z)(op).takeWhile(p).lastOption.getOrElse(z)
  }
}

by @ElBaulP (I modified a little bit to match comment by @Dima)

implicit class FoldLeftWhile[A](seq: Seq[A]) {

  def foldLeftWhile[B](z: B)(p: B => Boolean)(op: (B, A) => B): B = {
    @tailrec
    def foldLeftInternal(acc: B, seq: Seq[A]): B = seq match {
      case x :: _ =>
        val newAcc = op(acc, x)
        if (p(newAcc))
          foldLeftInternal(newAcc, seq.tail)
        else
          acc
      case _ => acc
    }

    foldLeftInternal(z, seq)
  }
}

Answer by me (involving side effects)

implicit class FoldLeftWhile[A](seq: Seq[A]) {

  def foldLeftWhile[B](z: B)(p: B => Boolean)(op: (B, A) => B): B = {
    var accumulator = z
    seq
      .map { e =>
        accumulator = op(accumulator, e)
        accumulator -> e
      }
      .takeWhile { case (acc, _) =>
        p(acc)
      }
      .lastOption
      .map { case (acc, _) =>
        acc
      }
      .getOrElse(z)
  }
}

回答6:

Simply use a branch condition on the accumulator:

seq.foldLeft(0, _ < 3) { (acc, e) => if (acc < 3) acc + e else acc}

However you will run every entry of the sequence.

来源：https://stackoverflow.com/questions/53728698/scala-foldleft-while-some-conditions-are-true

标签

scala

fold

foldleft