Efficiently traversing array matching a boolean condition?

Question

This is another question from Scala for the Impatient which says

Write a function lteqgt(values: Array[Int], v: Int) that returns a triple containing the counts of values less than v, equal to v, and greater than v.

The way I have done that is

scala> def lteqgt(values: Array[Int], v: Int): (Int, Int, Int) = (values.count(_ < v), values.count(_ == v), values.count(_ > v))
lteqgt: (values: Array[Int], v: Int)(Int, Int, Int)

scala> lteqgt(Array(0,0,1,1,1,2,2), 1)
res47: (Int, Int, Int) = (2,3,2)

Problem?
I am traversing the array 3 times to collect the counts, Is there a way to collect the values in first go? an idiomatic way?

Answer 1

This is perfect case to use foldLeft. It will traverse your collection exactly once, without creating another collection (as groupBy does) and it's more concise comparing to the more generic aggregate.

def lteqgt(values: Array[Int], v: Int): (Int, Int, Int) =
  values.foldLeft((0, 0, 0)) {
    case ((lt, eq, gt), el) =>
      if (el < v) (lt + 1, eq, gt)
      else if (el == v) (lt, eq + 1, gt)
      else (lt, eq, gt + 1)
  }

If you want to achieve ultimate efficiency, while avoiding imperative approach, then tail-recursion is the way to go:

def lteqgt(values: Array[Int], v: Int): (Int, Int, Int) = {
  def rec(i:Int, lt:Int, eq:Int, gt:Int):(Int, Int, Int) =
    if (i == values.length) (lt, eq, gt)
    else if (values(i) < v) rec(i + 1, lt + 1, eq, gt)
    else if (values(i) == v) rec(i + 1, lt, eq + 1, gt)
    else rec(i + 1, lt, eq, gt + 1)
  rec(0, 0, 0, 0)
}

This avoids construction Tuple and boxed Ints on each iteration. The whole thing compiles to while loop in java (here is the transpiled output if you are interested).

Answer 2

While a functional solution is arguably more elegant, let's not forget the "boring" but efficient imperative solution.

def lteqgt(values: Array[Int], v: Int): (Int, Int, Int) = {
    var lt = 0
    var eq = 0
    var gt = 0
    values.foreach (i => {
      if      (i<v)  lt += 1
      else if (i==v) eq += 1
      else           gt += 1
    })
    (lt,eq,gt)
  }

The above is likely to generate one function call per loop, as Aivean points out below. For more efficiency, we can remove the closure manually. (It's a pity that such optimization is not yet done by the compiler, though)

 def lteqgt(values: Array[Int], v: Int): (Int, Int, Int) = {
   var lt = 0
   var eq = 0
   var gt = 0
   var i = 0
   while (i < values.length) {
     if      (values(i) < v ) lt += 1
     else if (values(i) == v) eq += 1
     else                     gt += 1
     i += 1
   }
   (lt,eq,gt)
  }