I need to iterate over all K-tuples of nonnegative Ints that sum to N.
Eg if K is 3 and N==3,
(0,0,3)
(0,1,2)
(0,2,1)
(0,3,0)
(1,0,2)
(1,1,1)
(1,2,0)
(2,0,1)
(2,1,0)
(3,0,0)
Ordering does not matter, just full traversal. I need this to be allocation free though.
At the moment, I am not even sure of the algorithm that I would use, so any hints on anything between that and “package X does this” would be helpful.
do you need it specifically as an iterator? or would a foreach_tuple(f::Function, N, K) pattern be sufficient
well, my bid is
@generated function foreach_integer_simplex(f::F, N::Int, ::Val{K}) where {F, K}
if K == 1
return :(f((N,)))
end
syms = [Symbol(:i_, d) for d in (K-1):-1:1]
quote
s = 0
Base.Cartesian.@nloops $(K-1) i d -> 0:(N - s) d -> (s += i_d) d -> (s -= i_d) begin
remainder = N - s
val = ( $(syms...), remainder)
f(val)
end
end
end
julia> foreach_integer_simplex(3, Val{3}()) do t println(t) end
(0, 0, 3)
(0, 1, 2)
(0, 2, 1)
(0, 3, 0)
(1, 0, 2)
(1, 1, 1)
(1, 2, 0)
(2, 0, 1)
(2, 1, 0)
(3, 0, 0)
There are some algorithms here: https://jjj.de/fxt/fxtbook.pdf#chapter.16
And there’s an iterator in SmallCombinatorics.jl · SmallCombinatorics.jl which generates all the partitions, I suppose you can combine it with permutations in the same package, and filter out identitcal permutations. Or perhaps @matthias314 has something up his sleeve?
The function weakcompositions from SmallCombinatorics.jl gives what you want:
julia> using SmallCombinatorics, Chairmarks
julia> weakcompositions(3, 3) |> collect
10-element Vector{SmallVector{16, Int8}}:
[0, 0, 3]
[0, 1, 2]
[0, 2, 1]
[0, 3, 0]
[1, 0, 2]
[1, 1, 1]
[1, 2, 0]
[2, 0, 1]
[2, 1, 0]
[3, 0, 0]
julia> n = 3; k = 3; @b sum(first, weakcompositions($n, $k))
41.263 ns
The difference to @adienes’ solution is that the k parameter is not a
Val.
EDIT: Also, it’s an iterator, and it’s slower than @adienes’s code, see below.
also this 
julia> mutable struct Accumulator x::Int end
julia> add!(acc::Accumulator, y) = (acc.x += y)
add! (generic function with 1 method)
julia> const acc = Accumulator(0)
Accumulator(0)
julia> @b foreach_integer_simplex(Base.Fix1(add!, acc) ∘ sum, 8, Val{8}())
3.595 μs
julia> @b foreach(Base.Fix1(add!, acc) ∘ sum, weakcompositions(8, 8))
19.792 μs
another package which does the job:
julia> using Combinat
julia> collect(Compositions(3,3;min=0))
10-element Vector{Vector{Int64}}:
[0, 0, 3]
[0, 1, 2]
[0, 2, 1]
[0, 3, 0]
[1, 0, 2]
[1, 1, 1]
[1, 2, 0]
[2, 0, 1]
[2, 1, 0]
[3, 0, 0]
I think the OP is looking for an allocation-free solution. That would not be the case with Combinat.jl.
My idea is to enumerate starting from 0 and, after each increment by 1, check whether the sum of the digits equals N; if it does, I print it. This implementation has only a single allocation and runs very fast.
using StaticArrays
function odometer_inc_for!(d::MVector{K,Int}, ∑::Int) where {K}
@inbounds for i = 1:K
d[i] += 1
∑ += 1
if d[i] != 10
return ∑, false
end
d[i] = 0
∑ -= 10
end
return ∑, true # overflow
end
function scan_scheme1(::Val{K}, N::Integer) where {K}
d = zeros(MVector{K,Int})
∑ = 0
overflow = false
while !overflow
if ∑ == N
t = ntuple(i -> Int(d[K-i+1]), Val(K))
println(t)
end
∑, overflow = odometer_inc_for!(d, ∑)
end
nothing
end
scan_scheme1(Val(3), 3)
It doesn’t seem to be fast for larger inputs:
julia> const acc = Accumulator(0); # as before
julia> @b foreach_integer_simplex(Base.Fix1(add!, acc) ∘ first, 8, Val{8}())
6.769 μs
julia> @b sum(first, weakcompositions(8, 8))
20.865 μs
julia> @b scan_scheme2(Base.Fix1(add!, acc) ∘ first, Val{8}(), 8)
238.037 ms (1 allocs: 80 bytes, without a warmup)
Here I’m using a modification that takes a function as first argument:
function scan_scheme2(f::F, ::Val{K}, N::Integer) where {F,K}
d = zeros(MVector{K,Int})
∑ = 0
overflow = false
while !overflow
if ∑ == N
t = ntuple(i -> Int(d[K-i+1]), Val(K))
f(t)
end
∑, overflow = odometer_inc_for!(d, ∑)
end
nothing
end
Maybe you should clarify/relax this assumption.
First of all, your desired API probably needs to pass K as Val{K}(): Otherwise the return type of the iteration is not inferred, and you get dynamic dispatch and the tuple needs to be allocated itself.
One reasonable API is iterate_over_tuples(K,N) do the_tuple; ... end. This automatically introduces a function barrier to make the thing type stable. And it makes recursive implementations much simpler. But you probably can expect an allocation, due to the passing of an uninferred function barrier.
If you don’t want this, then K must be known at the callsite as a type. I.e. the user needs to have received ::Val{K} is its input already.
If you already explicitly construct / pass around ::Val{K} somewhere, why not have a small re-usable helper structure mutable struct IterationHelper{K, IntType} N::IntType ... end that is allocated but re-usable, and has some internal scratch space? That simplifies the task a lot!
I think “no allocations” is too hard of an ask; and plainly not what you actually need. A better ask would be:
O(K*log(N)), i.e. comparable to boxing the first output of the iterator, independent of the number of output tuples produced.That was implied, I have been using Julia for a decade now.
Not at all. After I slept on it, it turned out to be particularly simple:
"""
Generate the next item in the sequence of all `(i1, i2, ...)::NTuple{K,Int}` such that
`m ≥ i1 ≥ … ≥ 0`, ordered lexicographically using `>`.
NOTE: the invariant is not maintained if you start with `(m, m, …, m)`.
"""
function __inc(m::Int, i1::Int, iτ::Int...)
if i1 < m
(i1 + 1, iτ...)
else
iτ′ = __inc(m, iτ...)
(first(iτ′), iτ′...)
end
end
__inc(m::Int, i1::Int) = (i1 + 1,)
__diff(i1::Int, iτ::Int...) = (i1 - first(iτ), __diff(iτ...)...)
__diff(i1::Int) = i1
struct IntegerSimplex{K}
m::Int
function IntegerSimplex{K}(m::Int) where K
@assert K ≥ 1
@assert m ≥ 0
new{K}(m)
end
end
Base.length(itr::IntegerSimplex{K}) where K = binomial(itr.m + K - 1, K - 1)
Base.eltype(::IntegerSimplex{K}) where K = NTuple{K,Int}
function Base.iterate(itr::IntegerSimplex{K}, ι = ntuple(_ -> 0, Val(K-1))) where K
(; m) = itr
last(ι) > m && return nothing
ι′ = __inc(itr.m, ι...)
__diff(m, ι...), ι′
end
Eg
julia> @allocated mapreduce(sum, +, IntegerSimplex{3}(2))
0
Thanks for everyone who has answered! The key to my simple solution is to use a sequence m >= i1 >= ... internally, and then first difference it. This is much easier to keep track of when incrementing indices than working directly with the desired result (which is also doable, but I found it very confusing).
With SmallCombinatorics.jl the return type is inferred without Val(K) (and is some SmallVector).
That’s also what SmallCombinatorics.weakcompositions does internally.
Yes, I see, it is very clever design, basically restricted to a “large enough” SmallVector.
I think I will end up using that package, like SmallCollections.jl it looks very well designed. But figuring this out was fun
Iterators.product((0:10 for _= 1:3)...)