Scala: transform a collection, yielding 0..many elements on each iteration

Question

Given a collection in Scala, I'd like to traverse this collection and for each object I'd like to emit (yield) from 0 to multiple elements that should be joined together into a new collection.

For example, I expect something like this:

val input = Range(0, 15)
val output = input.somefancymapfunction((x) => {
  if (x % 3 == 0)
    yield(s"${x}/3")
  if (x % 5 == 0)
    yield(s"${x}/5")
})

to build an output collection that will contain

(0/3, 0/5, 3/3, 5/5, 6/3, 9/3, 10/5, 12/3)

Basically, I want a superset of what filter (1 → 0..1) and map (1 → 1) allows to do: mapping (1 → 0..n).

Solutions I've tried

Imperative solutions

Obviously, it's possible to do so in non-functional maneer, like:

var output = mutable.ListBuffer()
input.foreach((x) => {
  if (x % 3 == 0)
    output += s"${x}/3"
  if (x % 5 == 0)
    output += s"${x}/5"
})

Flatmap solutions

I know of flatMap, but it again, either:

1) becomes really ugly if we're talking about arbitrary number of output elements:

val output = input.flatMap((x) => {
  val v1 = if (x % 3 == 0) {
    Some(s"${x}/3")
  } else {
    None
  }
  val v2 = if (x % 5 == 0) {
    Some(s"${x}/5")
  } else {
    None
  }
  List(v1, v2).flatten
})

2) requires usage of mutable collections inside it:

val output = input.flatMap((x) => {
  val r = ListBuffer[String]()
  if (x % 3 == 0)
    r += s"${x}/3"
  if (x % 5 == 0)
    r += s"${x}/5"
  r
})

which is actually even worse that using mutable collection from the very beginning, or

3) requires major logic overhaul:

val output = input.flatMap((x) => {
  if (x % 3 == 0) {
    if (x % 5 == 0) {
      List(s"${x}/3", s"${x}/5")
    } else {
      List(s"${x}/3")
    }
  } else if (x % 5 == 0) {
    List(s"${x}/5")
  } else {
    List()
  }
})

which is, IMHO, also looks ugly and requires duplicating the generating code.

Roll-your-own-map-function

Last, but not least, I can roll my own function of that kind:

def myMultiOutputMap[T, R](coll: TraversableOnce[T], func: (T, ListBuffer[R]) => Unit): List[R] = {
  val out = ListBuffer[R]()
  coll.foreach((x) => func.apply(x, out))
  out.toList
}

which can be used almost like I want:

val output = myMultiOutputMap[Int, String](input, (x, out) => {
  if (x % 3 == 0)
    out += s"${x}/3"
  if (x % 5 == 0)
    out += s"${x}/5"
})

Am I really overlooking something and there's no such functionality in standard Scala collection libraries?

Similar questions

This question bears some similarity to Can I yield or map one element into many in Scala? — but that question discusses 1 element → 3 elements mapping, and I want 1 element → arbitrary number of elements mapping.

Final note

Please note that this is not the question about division / divisors, such conditions are included purely for illustrative purposes.

Answer 1

Rather than having a separate case for each divisor, put them in a container and iterate over them in a for comprehension:

val output = for {
  n <- input
  d <- Seq(3, 5)
  if n % d == 0
} yield s"$n/$d"

Or equivalently in a collect nested in a flatMap:

val output = input.flatMap { n =>
  Seq(3, 5).collect {
    case d if n % d == 0 => s"$n/$d"
  }
}

In the more general case where the different cases may have different logic, you can put each case in a separate partial function and iterate over the partial functions:

val output = for {
  n <- input
  f <- Seq[PartialFunction[Int, String]](
    {case x if x % 3 == 0 => s"$x/3"},
    {case x if x % 5 == 0 => s"$x/5"})
  if f.isDefinedAt(n)
} yield f(n)

Answer 2

You can also use some functional library (e.g. scalaz) to express this:

import scalaz._, Scalaz._

def divisibleBy(byWhat: Int)(what: Int): List[String] = 
  (what % byWhat == 0).option(s"$what/$byWhat").toList

(0 to 15) flatMap (divisibleBy(3) _ |+| divisibleBy(5))

This uses the semigroup append operation |+|. For Lists this operation means a simple list concatenation. So for functions Int => List[String], this append operation will produce a function that runs both functions and appends their results.

Answer 3

If you have complex computation, during which you should sometimes add some elements to operation global accumulator, you can use popular approach named Writer Monad

Preparation in scala is somewhat bulky but results are extremely composable thanks to Monad interface

import scalaz.Writer
import scalaz.syntax.writer._
import scalaz.syntax.monad._
import scalaz.std.vector._
import scalaz.syntax.traverse._

type Messages[T] = Writer[Vector[String], T]

def yieldW(a: String): Messages[Unit] = Vector(a).tell

val output = Vector.range(0, 15).traverse { n =>
  yieldW(s"$n / 3").whenM(n % 3 == 0) >>
  yieldW(s"$n / 5").whenM(n % 5 == 0)
}.run._1 

Answer 4

Here is my proposition for a custom function, might be better with pimp my library pattern

def fancyMap[A, B](list: TraversableOnce[A])(fs: (A => Boolean, A => B)*) = {
  def valuesForElement(elem: A) = fs collect { case (predicate, mapper) if predicate(elem) => mapper(elem) }
  list flatMap valuesForElement
}

fancyMap[Int, String](0 to 15)((_ % 3 == 0, _ + "/3"), (_ % 5 == 0, _ + "/5"))

Answer 5

You can try collect:

val input = Range(0,15)
val output = input.flatMap { x =>
     List(3,5) collect { case n if (x%n == 0) => s"${x}/${n}" }
}
System.out.println(output)

Answer 6

I would us a fold:

val input = Range(0, 15)
val output = input.foldLeft(List[String]()) {
    case (acc, value) =>
        val acc1 = if (value % 3 == 0) s"$value/3" :: acc else acc
        val acc2 = if (value % 5 == 0) s"$value/5" :: acc1 else acc1
        acc2
}.reverse

output contains

List(0/3, 0/5, 3/3, 5/5, 6/3, 9/3, 10/5, 12/3)

A fold takes an accumumlator (acc), a collection, and a function. The function is called with the initial value of the accumumator, in this case an empty List[String], and each value of the collection. The function should return an updated collection.

On each iteration, we take the growing accumulator and, if the inside if statements are true, prepend the calculation to the new accumulator. The function finally returns the updated accumulator.

When the fold is done, it returns the final accumulator, but unfortunately, it is in reverse order. We simply reverse the accumulator with .reverse.

There is a nice paper on folds: A tutorial on the universality and expressiveness of fold, by Graham Hutton.