scalasequencefor-comprehensionresulttype

Scala: Iterate an Array's indices in for-comprehension but ensure result type is same Array type

I know that the sequence type of the first sequence found in a for-comprehension defines the output type of the comprehension. But I need to get around this without sacrificing the use of comprehension syntax.

Suppose I have some Array[Double] called v and some complicated predicate function called condition that is based on the indices of v, and if condition(i) evaluates to true, then I want to keep the element of v. It's similar to using filter, except the filtering happens on the indices, not on the values of v.

I want to express this efficiently with a for-comprehension over the indices, but I want the type of the result to be Array[Double] and not whatever is the type of Array.indices.

Here's a minimal example where condition is just a toy example of calculating whether an index is even. Real use cases involve checking indices in a 2D array where certain image gradient conditions are true, but the idea is the same as this simple example.

scala> val v = Array[Double](8.0, 9.0, 10.0, 11.0, 12.0)
v: Array[Double] = Array(8.0, 9.0, 10.0, 11.0, 12.0)

scala> def condition(x: Int): Boolean = {x % 2 == 0} // but could be complex
condition: (x: Int)Boolean

scala> for (i <- v.indices if condition(i)) yield v(i)
res136: scala.collection.immutable.IndexedSeq[Double] = Vector(8.0, 10.0, 12.0)

The output of the for-comprehension is of type scala.collection.immutable.IndexedSeq[Double] but I am seeking how to ensure it will just be Array[Double], without needing an explicit type cast.

Basically, if I iterate from foo.indices, and the items yielded come from foo(i), I am seeking a way to automatically use the type of foo for the result, so that the container type matches whatever type's indices were iterated over.

What is the right idiom in Scala for ensuring the result type of a for-comprehension that needs to specifically range over indices, but where the container type of the result should match the container type of the thing-whose-indices-are-being-used instead of the type of the indices sequence itself (which is inconsequential for the comprehension)?

Solution

    1. Use zipWithIndex. In case of Array, it will produce Array[(Double, Int)], not some strange IndexedSeq:

      for ((a, i) <- v.zipWithIndex if condition(i)) yield a

    2. Use breakOut

      import collection.breakOut
val res: Array[Double] = (
  for (i <- v.indices if condition(i)) yield v(i)
)(breakOut)

    3. Just convert the result in the end with toArray, optimize it later when it turns out to be really necessary.

    4. Also note that zipWithIndex can be defined in general for all collections that have a Traverse typeclass. It is a general method to convert F[A] into an F[(A, Int)] under assumption that there is a Traverse[F] available.