Abstract over repeated flatMap

Question

I am trying to generalize repeated, nested flatMap but not sure if one exists.

The following code will produce all combinations of n choose 3, :

def choose3flatMap(n: Int, r: Int = 3) =
  (0 to n - r)
    .flatMap(i => (i + 1 to n - (r - 1))
      .flatMap(j => (j + 1 to n - (r - 2))
        .map(k => Seq(i, j, k))))

Repeating the flatMap operation, we can get all combinations of n choose 5, :

def choose5flatMap(n: Int, r: Int = 5) =
  (0 to n - r)
    .flatMap(i => (i + 1 to n - (r - 1))
      .flatMap(j => (j + 1 to n - (r - 2))
        .flatMap(k => (k + 1 to n - (r - 3))
          .flatMap(l => (l + 1 to n - (r - 4))
            .map(m => Seq(i, j, k, l, m)))))

Clearly there is a pattern here. I would like to utilize this similarity to get a general solution for n choose r, . Is there a simple way to accomplish this. Perhaps a higher order function of some sort?

What I have tried:

Scala lets me rewrite the map/flatMap with a for expression. This reads cleaner, but the number of choices in still hard-coded.

def choose3Loop(n: Int, r: Int = 3) =
  for {
    i <- 0 to n - r
    j <- i + 1 to n - (r - 1)
    k <- j + 1 to n - (r - 2)
  } yield Seq(i, j, k)

I can write a recursive solution directly using flatMap or utilizing the sugar of a for expression:

def combinationsRecursive(n: Int, r: Int, i: Int = 0): Seq[Seq[Int]] =
  if (r == 1) (i until n).map(Seq(_))
  else {
    (i to n - r).flatMap(
      i => combinationsRecursive(n, r - 1, i + 1).map(j => i +: j))
  }

def combinationsRecursiveLoop(n: Int, r: Int, i: Int = 0): Seq[Seq[Int]] =
  if (r == 1) (i until n).map(Seq(_))
  else
    for {
      i <- i to n - r
      j <- combinationsRecursiveLoop(n, r - 1, i + 1)
    } yield i +: j

While these are solutions to the general problem, I wonder if there is a higher-level abstraction I am missing here that may be applicable to other problems as well. I recognize that for this particular application, I could do (0 to n).combinations(r) to use a library-provided implementation of computing combinations.

While the above code is Scala, in this case I am interested the functional programming aspect of it and not the language capabilities. If there is a solution but one that is not supported by Scala I am interested in that.

Edit: He is a sample caller and the resulting output by request:

scala> combinationsRecursiveLoop(5, 3)

res0: Seq[Seq[Int]] = Vector(List(0, 1, 2), List(0, 1, 3), List(0, 1, 4), List(0, 2, 3), List(0, 2, 4), List(0, 3, 4), List(1, 2, 3), List(1, 2, 4), List(1, 3, 4), List(2, 3, 4))

scala> combinationsRecursiveLoop(5, 3).map("("+_.mkString(", ")+")").mkString(" ")

res1: String = (0, 1, 2) (0, 1, 3) (0, 1, 4) (0, 2, 3) (0, 2, 4) (0, 3, 4) (1, 2, 3) (1, 2, 4) (1, 3, 4) (2, 3, 4)

It just provides all r-element subsets of the set of integers starting at zero containing n elements. More information on combinations can be found on Wikipedia.

Answer 1

Here is one way to look at this, that I have come up with.

You can extract one stage in your chain as a function f: List[Int] => List[List[Int]], that takes a List with a beginning of a combination, and prepends all possible next elements to it.

For example in choose(5, 3), f(List(2, 0)) would result in List(List(3, 2, 0), List(4, 2, 0)).

Here is a possible implementation of such a function with some processing for the initial case added:

val f: List[Int] => List[List[Int]] = l =>
  (l.headOption.map(_ + 1).getOrElse(0) to n - (r - l.size))
    .map(_ :: l).toList

Now, such a function is a Kleisli arrow Kleisli[List, List[Int], List[Int]], and it's endomorphic (has the same argument and return types).

There is a monoid instance for endomorphic kleisli arrows, where the monoid "addition" means the flatMap operation (or in pseudocode, f1 |+| f2 == a => f1(a).flatMap(f2)). So to replace your chain of flatMaps you need to "add" r instances of this f function, or in other words to multiply the f function by r.

This idea translates directly into Scalaz code:

import scalaz._, Scalaz._

def choose(n: Int, r: Int) = {
  val f: List[Int] => List[List[Int]] = l =>
    (l.headOption.map(_ + 1).getOrElse(0) to n - (r - l.size))
      .map(_ :: l).toList
  Endomorphic.endoKleisli(f).multiply(r).run(Nil)
}

And here is an example running it:

scala> choose(4, 3)
res1: List[List[Int]] = List(List(2, 1, 0), List(3, 1, 0), List(3, 2, 0), List(3, 2, 1))

The combinations are reversed, but it should be possible to make a version, that produces combinations with elements in the increasing order (or just run choose(n, r).map(_.reverse)).

Another improvement would be to make a lazy version, that returns Stream[List[Int]] (or even better a scalaz.EphemeralStream[List[Int]]: you don't want to have all the combinations cached in memory), but this is left as an exercise to the reader.