How to find all possible combinations of elements of multiple lists?

Question

I need to find all combinations of multiple lists contained in a list of lists. For example:

Given [[1, 2, 3], [4, 5], [6]], I need to obtain (order not important)

[[1, 4, 6], [1, 5, 6], [2, 4, 6], [2, 5, 6], [3, 4, 6], [3, 5, 6]].

If I knew ahead of time how many lists were contained in the original list of lists (and this number were sufficiently small), it would be easy enough to do nested for loops to generate the combinations. But my situation is that I do not know this ahead of time. How would I go about generating this list?

Although I am using Swift, I suppose this problem is more general than that. If there is a specific Swift method which does this that would work for me.

Answer 1

You can find out the expected number of combinations by multiplying the sizes of the sub-arrays. So [[1, 2, 3], [4, 5], [6]] will generate 3 * 2 * 1 = 6 combinations.

Here is a generic function that will combine multiple arrays to create an array of all combinations of elements from each of the original sub-arrays.

It works recursively by building up partial lists by removing the first sub-array from the original array and then adding those elements to the partial list before calling itself recursively.

When combine is eventually called with an empty list, it returns the partial list which is now complete.

The returned lists get collected in the variable result and returned from combine.

// combine elements from the sub-arrays of lists
// partial holds the partial results as the combinations are built up
// partial has a default value of empty array for the initial call

func combine<T>(lists: [[T]], partial: [T] = []) -> [[T]] {
    // print("combine(lists: \(lists), partial: \(partial))")
    if lists.isEmpty {
        // recursive base case: lists is now empty, so partial
        // is complete, so return it in an enclosing array
        // print("... returning \([partial])")
        return [partial]
    } else {
        // make lists mutable so that we can remove the first sub-array
        var lists = lists

        // remove the first sub-array from lists which is now shorter
        let first = lists.removeFirst()

        // create an array to hold all of the combinations
        var result = [[T]]()

        // take each element from the first sub-array, append it to
        // the partial result, and call combine to continue the
        // process.  Take the results returned from combine and append
        // those to the result array.
        for n in first {
            result += combine(lists: lists, partial: partial + [n])
        }

        // Return the results
        // print("... returning \(result)")
        return result
    }
}

Test:

let result = combine(lists: [[1, 2, 3], [4, 5], [6]])
print(result)

Output:

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

---

If you uncomment the three print statements, you can get some insight how this works:

combine(lists: [[1, 2, 3], [4, 5], [6]], partial: [])
combine(lists: [[4, 5], [6]], partial: [1])
combine(lists: [[6]], partial: [1, 4])
combine(lists: [], partial: [1, 4, 6])
... returning [[1, 4, 6]]
... returning [[1, 4, 6]]
combine(lists: [[6]], partial: [1, 5])
combine(lists: [], partial: [1, 5, 6])
... returning [[1, 5, 6]]
... returning [[1, 5, 6]]
... returning [[1, 4, 6], [1, 5, 6]]
combine(lists: [[4, 5], [6]], partial: [2])
combine(lists: [[6]], partial: [2, 4])
combine(lists: [], partial: [2, 4, 6])
... returning [[2, 4, 6]]
... returning [[2, 4, 6]]
combine(lists: [[6]], partial: [2, 5])
combine(lists: [], partial: [2, 5, 6])
... returning [[2, 5, 6]]
... returning [[2, 5, 6]]
... returning [[2, 4, 6], [2, 5, 6]]
combine(lists: [[4, 5], [6]], partial: [3])
combine(lists: [[6]], partial: [3, 4])
combine(lists: [], partial: [3, 4, 6])
... returning [[3, 4, 6]]
... returning [[3, 4, 6]]
combine(lists: [[6]], partial: [3, 5])
combine(lists: [], partial: [3, 5, 6])
... returning [[3, 5, 6]]
... returning [[3, 5, 6]]
... returning [[3, 4, 6], [3, 5, 6]]
... returning [[1, 4, 6], [1, 5, 6], [2, 4, 6], [2, 5, 6], [3, 4, 6], [3, 5, 6]]