Type-stable alternatives to Iterators.product?

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

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

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