How to create lazy combinations

Question

My question is very simple, how do I make this code lazy:

/*
input: [
    [1, 2],
    [3, 4],
    [5, 6]
]

output: [
    [1, 3, 5],
    [1, 3, 6],
    [1, 4         


        

Answer 1

Here is what I came up with:

func combinations<T>(options: [[T]]) -> AnySequence<[T]> {
    guard let lastOption = options.last else {
        return AnySequence(CollectionOfOne([]))
    }
    let headCombinations = combinations(options: Array(options.dropLast()))
    return AnySequence(headCombinations.lazy.flatMap { head in
        lastOption.lazy.map { head + [$0] }
    })
}

The main difference to this solution is that the recursive call creates a sequence of the first N-1 options, and then combines each element of that sequence with each element of the last option. This is more efficient because the sequence returned from the recursive call is enumerated only once, and not once for each element that it is combined with.

Other differences are:

• There is no need to call .lazy on the AnySequence if that sequence is already lazy. The return type is therefore "simplified" to AnySequence<[T]>.
• I have used CollectionOfOne to create a single-element sequence for the empty array.
• Treating the case options.count == 1 separately is not necessary for the algorithm to work (but might be a possible performance improvement).

---

A completely different approach is to define a custom collection type which computes each combination as a function of the index, using simple modulo arithmetic:

struct Combinations<T> : RandomAccessCollection {
    let options: [[T]]
    let startIndex = 0
    let endIndex: Int

    init(options: [[T]]) {
        self.options = options.reversed()
        self.endIndex = options.reduce(1) { $0 * $1.count }
    }

    subscript(index: Int) -> [T] {
        var i = index
        var combination: [T] = []
        combination.reserveCapacity(options.count)
        options.forEach { option in
            combination.append(option[i % option.count])
            i /= option.count
        }
        return combination.reversed()
    }
}

No extra storage is needed and no recursion. Example usage:

let all = Combinations(options: [[1, 2], [3, 4], [5, 6]])
print(all.count)
for c in all { print(c) }

Output:

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

Testing with

let options = Array(repeating: [1, 2, 3, 4, 5], count: 5)

this collection-based method turned out to be faster then the my above sequence-based method by a factor of 2.

Answer 2

I found one possible solution, but I'll leave this answer not accepted for a while to see if someone knows a better one.

func combinations<T>(options: [[T]]) -> LazySequence<AnySequence<[T]>> {
    guard let head = options.first else {
        return AnySequence([].lazy.map({ [$0] })).lazy
    }

    if options.count == 1 {
        return AnySequence(head.lazy.map({ [$0] })).lazy
    }

    let tailCombinations = combinations(options: Array(options.dropFirst()))

    return AnySequence(head.lazy.flatMap({ option in
        return tailCombinations.lazy.map({ [option] + $0 })
    })).lazy
}

The solution was to use AnySequence instead of AnyCollection. I'm not sure why though, I'd still like to have the AnyCollection interface rather than AnySequence, since it provides me with a few more methods, like count.