Type-stable alternatives to Iterators.product?

General Usage
,
1

I’m dealing with a combinatorial problem. Given N vectors of integers of varied lengths, I want to find all possible ways to select one integer from each vector. Here’s my code, which is unfortunately not type-stable:

julia> const data = [[1,2], [11,12], [101]];

julia> function all_combinations(a::Vector{Vector{Int}})
           combinations = Vector{Int}[]
           for c in Iterators.product(a...)
               push!(combinations, collect(c))
           end
           combinations
       end

julia> all_combinations(data)
4-element Vector{Vector{Int64}}:
 [1, 11, 101]
 [2, 11, 101]
 [1, 12, 101]
 [2, 12, 101]

Type instability is present according to the output of @code_warntype all_combinations(data). Is there a neat way to rewrite the code in a type-stable manner? (Neither N nor the length of each vector is known at compile time.)

1 Like
2

This works

function all_combinations(a)
    combinations = Vector{Int}[]
    for c in Iterators.product(a...)
        push!(combinations, collect(c))
    end
    combinations
end
data = ([1,2], [11,12], [101])

@code_warntype all_combinations(data)

The trick is that the a... operation is only type stable for things with known types and length!
One way to make the length to be known at compile time is to use a Tuple.

2 Likes
4

I’ve come up with a type-stable solution when the input has unknown length:

function all_combinations(a::Vector{Vector{Int}})
    foldl(form_combinations, a, init = Vector{Int}[])
end

function form_combinations(a::Vector{Vector{Int}}, b::Vector{Int})
    if length(a) == 0
        [[i] for i in b]
    else
        [push!(copy(i), j) for i in a for j in b]
    end
end

Test:

julia> all_combinations([[1,2], [11, 12], [101]])
4-element Vector{Vector{Int64}}:
 [1, 11, 101]
 [1, 12, 101]
 [2, 11, 101]
 [2, 12, 101]
2 Likes