Finding all arrays of length L that add up to a certain number

I want to create an array of tuples where each tuple is of length L and the sum of all elements come to a defined number.

I do have a solution currently, but it is not scalable till what I need, I get thrown an OutOfMemoryError().

Current solution:

function calc_combinations(C, n_var)
    x = 0:C
    comb = []
    for i in Iterators.product(fill(x, n_var)...) |> collect
        if sum(i) == C && !(i in comb)
            push!(comb, i)
        end
    end
    return comb
end

For larger numbers, for both n_var and C, I get thrown the error at Iterators.product(fill(x, n_var)…) |> collect.

For my purpose, I only need to go till C=20 and n_var=8 at max (memory runs out at this point as well), but how would one go about making it much more efficient and also scalable as well, or if there is any other way to go about solving this problem.

Don’t collect the iterator, i.e., the whole point of an iterator is that you can run through it without allocating a data structure holding all its elements.
The following seems not to run out of memory:

function calc_combinations(C, n_var)
    x = 0:C
    comb = Set()  # Use a set as we don't want duplicates (no need to check i in C!)
    for i in Iterators.product(fill(x, n_var)...)
        if sum(i) == C
            push!(comb, i)
        end
    end
    return comb
end

Also some further ideas:

  1. The brute-force solution can be written as a one-liner
    calc_combs(C, n_var) = collect(Iterators.filter(==(C)∘sum, Iterators.product(fill(0:C, n_vars)...)))
    
  2. For C = 20 and n_{var} = 8 there are 21^8 possible tuples, so this gets rather slow.
    A better solution would be to not generate all of these, but extend only potential candidate solutions:
    function calc_combs_faster(C, n_vars)
        # Extend a potential candidate, i.e., sum still below C
        extend(candidate) = [(candidate..., i) for i in 0:(C-sum(candidate; init=0))]
        candidates = [()]
        for _ in 1:n_vars
            candidates = reduce(vcat, extend.(candidates))
        end
        # Finally, only keep actual solutions
        filter(==(C)∘sum, candidates)
    end
    

I wrote this function a long time ago in Combinatorics.jl, it is called multiexponents:

Thanks @juliohm and @bertschi for the answers, all these implementations worked. I do need to try and understand how iterators work properly it seems. But all these implementations are elegant and quick.

Thanks a lot!

I think what you want is a composition:

I submitted a PR a while ago (probably stale now) to add this capability to Combinatorics.jl:

It provides an iterator like @bertschi suggested. But, compositions are returned as Vectors, not tuples.

The package Combinat contains a very fast implementation of compositions

julia> @btime compositions(5,3;min=0)
  758.054 ns (2 allocations: 1.44 KiB)
21-element Vector{SubArray{Int64, 1, Matrix{Int64}, Tuple{Int64, Base.Slice{Base.OneTo{Int64}}}, true}}:
 [0, 0, 5]
 [0, 1, 4]
 [0, 2, 3]
 [0, 3, 2]
 [0, 4, 1]
 [0, 5, 0]
 [1, 0, 4]
 [1, 1, 3]
 [1, 2, 2]
 [1, 3, 1]
 [1, 4, 0]
 [2, 0, 3]
 [2, 1, 2]
 [2, 2, 1]
 [2, 3, 0]
 [3, 0, 2]
 [3, 1, 1]
 [3, 2, 0]
 [4, 0, 1]
 [4, 1, 0]
 [5, 0, 0]
julia> ncompositions(20,8;min=0)
888030
ulia> @btime compositions(20,8;min=0);
  78.271 ms (4 allocations: 88.08 MiB)


2 Likes

If speed matters, then the code below seems to be even faster:

julia> n = 20; k = 8;

julia> @b compositions($n, $k; min=0)
63.224 ms (11 allocs: 88.077 MiB, 7.46% gc time)

julia> @b collect(Compositions{$k}($n))   # new
15.721 ms (10 allocs: 54.201 MiB, 11.82% gc time)

julia> @b sum(sum, compositions($n, $k; min=0))
74.697 ms (4 allocs: 88.077 MiB, 6.59% gc time)

julia> @b sum(sum, Compositions{$k}($n))   # new
1.875 ms (13 allocs: 272 bytes)
Code
import Base: iterate, length, eltype

struct Compositions{K}
    n::Int
end

eltype(::Type{Compositions{K}}) where K = NTuple{K, Int}

length(c::Compositions{K}) where K = binomial(c.n+K-1, c.n)

iterate(c::Compositions{1}, ::Any) = nothing

function iterate(c::Compositions{K}) where K
    t = ntuple(i -> i == 1 ? c.n : 0, K)
    cc = ntuple(Returns(0), Val(K-1))
    t, (cc, t)
end

@inline function iterate(c::Compositions{K}, (cc, t)) where K
    r = iterate(Compositions{K-1}(cc[1]), (cc[2:end], t[2:end]))
    if r !== nothing
        t1 = t[1]
        c1 = cc[1]
        _, (ncc, nt) = r
    elseif t[1] != 0
        t1 = t[1]-1
        c1 = c.n-t1
        _, (ncc, nt) = iterate(Compositions{K-1}(c1))
    else
        return nothing
    end
    return (t1, nt...), ((c1, ncc...), (t1, nt...))
end
2 Likes