Abstract over repeated flatMap - scala

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.