Search code examples
kotlinsetcartesian-productidioms

Idiomatic way to create n-ary cartesian product (combinations of several sets of parameters)


To create all possible combinations of two sets of parameters and perform an action on them, you can do:

setOf(foo, bar, baz).forEach { a ->
    setOf(0, 1).forEach { b ->
        /* use a and b */
    }
}

However, if you have (potentially many) more parameters, this quickly turns into a pyramid of doom:

setOf(foo, bar, baz).forEach { a ->
    setOf(0, 1).forEach { b ->
        setOf(true, false, null).forEach { c ->
            setOf("Hello,", "World!").forEach { d ->
                /* use a, b, c and d */
            }
        }
    }
}

You could write this similarly with for loops, or differently like so:

val dAction = { d: String -> /* use a, b, c and d */ }
val cAction = { c: Boolean? -> setOf("Hello,", "World!").forEach(dAction) }
val bAction = { b: Int -> setOf(true, false, null).forEach(cAction) }
val aAction = { a: Any? -> setOf(0, 1).forEach(bAction) }
setOf(foo, bar, baz).forEach(aAction)

But I don't think that's any better, because there are some readability issues here: d, c, b and a's actions are written in reverse. Their type specifications cannot be inferred, so they must be specified. It's reversed sequentially compared to the pyramid of doom. The order of the sets providing the possible values should not matter, but it does: you just want to create any combinations from a bunch of sets, however, in this code every line depends on the previous.

It would be very nice to have an idiomatic way of doing something like Python's or Haskell's comprehensions, in which you (almost like the mathematical notation) can do something like:

{ /* use a, b, c and d */
    for a in setOf(foo, bar, baz),
    for b in setOf(0, 1),
    for c in setOf(true, false, null),
    for d in setOf("Hello,", "World!")
}

Which is very easy to read: there is no excessive indentation, the action you're interested in goes first, the data sources are very clearly defined, etc.

Side note: similar problems occur with flatMap-flatMap-...-flatMap-map.

Any ideas about how to neatly create n-ary cartesian products in Kotlin?


Solution

  • I've created a solution myself, so I don't have to add a dependency as suggested by Omar's answer.

    I created a function that takes two or more sets of any size:

    fun cartesianProduct(a: Set<*>, b: Set<*>, vararg sets: Set<*>): Set<List<*>> =
        (setOf(a, b).plus(sets))
            .fold(listOf(listOf<Any?>())) { acc, set ->
                acc.flatMap { list -> set.map { element -> list + element } }
            }
            .toSet()
    

    Example:

    val a = setOf(1, 2)
    val b = setOf(3, 4)
    val c = setOf(5)
    val d = setOf(6, 7, 8)
    
    val abcd: Set<List<*>> = cartesianProduct(a, b, c, d)
    
    println(abcd)
    

    Output:

    [[1, 3, 5, 6], [1, 3, 5, 7], [1, 3, 5, 8], [1, 4, 5, 6], [1, 4, 5, 7], [1, 4, 5, 8], [2, 3, 5, 6], [2, 3, 5, 7], [2, 3, 5, 8], [2, 4, 5, 6], [2, 4, 5, 7], [2, 4, 5, 8]]
    

    The function cartesianProduct returns a set of lists. There's a number of problems with these lists:

    • Any type information is lost, because the returned set contains lists that contain the union of types of the input sets. The returned type of these lists' elements is Any?. The function returns a Set<List<*>>, i.e. Set<List<Any?>>.
    • By definition, the size of the lists isn't known; they are not like a Kotlin Pair or Triple, where the size is a constant by definition. However, the size of these lists/tuples should be equal to the number of input sets, i.e. 4 in the example above.

    However, using reflection, we can solve these problems. The action we want to take with every list can be written as a function (e.g. a constructor of some class, which is also just a function):

    data class Parameters(val number: Int, val maybe: Boolean?) {
        override fun toString() = "number = $number, maybe = $maybe"
    }
    
    val e: Set<Int> = setOf(1, 2)
    val f: Set<Boolean?> = setOf(true, false, null)
    
    val parametersList: List<Parameters> = cartesianProduct(e, f).map { ::Parameters.call(*it.toTypedArray()) }
    
    println(parametersList.joinToString("\n"))
    

    Output:

    number = 1, maybe = true
    number = 1, maybe = false
    number = 1, maybe = null
    number = 2, maybe = true
    number = 2, maybe = false
    number = 2, maybe = null
    

    The signature of the transform (::Parameters in the example) specifies the contract for the lists' contents.

    Because map { ::Parameters.call(*it.toTypedArray()) } is not very nice, I've created a second extension function that does it for me:

    fun <T> Set<List<*>>.map(transform: KFunction<T>) = map { transform.call(*it.toTypedArray()) }
    

    With that, the code becomes quite idiomatic:

    val parametersList: List<Parameters> = cartesianProduct(e, f).map(::Parameters)
    

    The code is available from this GitHub Gist, where I will update it if I ever improve it. There are also tests: the cartesian product that includes any empty set returns the empty set, as is mathematically expected. I'm neither saying that this is an optimal solution, nor that it is mathematically sound (not every mathematical property is explicitly implemented and tested), but it works for the question's purpose.