Iterating over integer simplex

Specific Domains Numerics
, ,
1

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.

1 Like
2

do you need it specifically as an iterator? or would a foreach_tuple(f::Function, N, K) pattern be sufficient

3

CartesianIndices to the rescue!
I love that one can use these to rock stuff. Per dimension you can provide a range. To have the numbers different I chose N=2 to see the difference and then you can do

 K=3; N=2; for x in CartesianIndices(Tuple(fill(0:N, K)))
       println(x.I)
 end

and get

(0, 0, 0)
(1, 0, 0)
(2, 0, 0)
(0, 1, 0)
(1, 1, 0)
(2, 1, 0)
(0, 2, 0)
(1, 2, 0)
(2, 2, 0)
(0, 0, 1)
(1, 0, 1)
(2, 0, 1)
(0, 1, 1)
(1, 1, 1)
(2, 1, 1)
(0, 2, 1)
(1, 2, 1)
(2, 2, 1)
(0, 0, 2)
(1, 0, 2)
(2, 0, 2)
(0, 1, 2)
(1, 1, 2)
(2, 1, 2)
(0, 2, 2)
(1, 2, 2)
(2, 2, 2)

It only allocates the tuple of ranges (from the fill)

edit: Of sorry I forgot the sum. Maybe this can still be used as a starting point.

4

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)
5

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?

6

The function
weakcompositions
from SmallCombinatorics.jl gives what you want:

> using SmallCombinatorics, Chairmarks

> 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]

> 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.

7

also this :wink:

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