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:

1 Like

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

Nice implementation. But do you have a variant where the min can be specified? Usually, in the literature, compositions are defined where the minimum of an entry is 1 (this way there is a finite number of compositions whatever the value of k).

To have a minimum of m, one could iterate over

compositions_min(n, k, m) = (t .+ m for t in Compositions{k}(n-k*m))

(As before, this assumes that there is at least one such composition.)

julia> compositions_min(6, 3, 1) |> collect
10-element Vector{Tuple{Int64, Int64, Int64}}:
 (4, 1, 1)
 (3, 2, 1)
 (3, 1, 2)
 (2, 3, 1)
 (2, 2, 2)
 (2, 1, 3)
 (1, 4, 1)
 (1, 3, 2)
 (1, 2, 3)
 (1, 1, 4)

Alternatively, one can modify Compositions. This is below, with some polishing of the code. Iā€™ve done a few experiments, and this change doesnā€™t seem to have a significant impact on the runtime.

Modified code
import Base: iterate, length, eltype

struct Compositions2{K}
    n::Int
    min::Int
end

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

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

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

function iterate(c::Compositions2{K}) where K
    ns = ntuple(Returns(c.min), Val(K-1))
    xs = (c.n-(K-1)*c.min, ns...)
    xs, (ns, xs)
end

@inline function iterate(c::Compositions2{K}, (ns, xs)) where K
    n1, nr... = ns
    x1, xr... = xs
    ret = iterate(Compositions2{K-1}(n1, c.min), (nr, xr))
    if ret !== nothing
        _, (ns2, xs2) = ret
    elseif x1 > c.min
        x1 -= 1
        n1 = c.n-x1
        _, (ns2, xs2) = iterate(Compositions2{K-1}(n1, c.min))
    else
        return nothing
    end
    new_xs = (x1, xs2...)
    new_ns = (n1, ns2...)
    return new_xs, (new_ns, new_xs)
end
1 Like

You could make a PR to Combinatorics with your last versionā€¦

@ararslan @logankilpatrick Would you be interested in this for Combinatorics.jl? Here is a comparison (with Julia 1.11.0-rc4):

julia> @b multiexponents(16, 8) sum(sum, _)
70.247 ms (2941886 allocs: 179.558 MiB, 22.75% gc time)

julia> @b Compositions2{8}(16, 0) sum(sum, _)
326.503 Ī¼s

Note that at present the new iterator returns tuples, not vectors.

1 Like