Custom iterator, simple example

General Usage
1

I implemented Donald Knuth’s “Algorithm M” for generating all numbers in a mixed-radix numbering system. I would like to turn this into an iterator so I could loop over all of them but without having them all in memory at the same time.

Here’s a simple case with the code:

m = [1,2,1]
nl = length(m)
a = zeros(Int,nl)
while true
    j = nl
    println(a)
    while true
      if a[j] < m[j] break; end
      a[j] = 0
      j -= 1
      if j==0 break; end
    end
    if j == 0 break; end
    a[j] += 1    
end

What I would like is something like this:

for i in MyMixedRadix([1,2,1])
   # code goes here
end

How do I turn my nested for loops into an iterator?

PS Here’s the output from the example:

[0, 0, 0]
[0, 0, 1]
[0, 1, 0]
[0, 1, 1]
[0, 2, 0]
[0, 2, 1]
[1, 0, 0]
[1, 0, 1]
[1, 1, 0]
[1, 1, 1]
[1, 2, 0]
[1, 2, 1]
1 Like
2

I can give you a start:

struct MixedRadix{T}
  m::T
end
function Base.iterate(mr::MixedRadix)
    nl = length(mr.m)
    a = zeros(Int, nl)
    a, a # first iterate and the initial state
end
function Base.iterate(mr::MixedRadix, a)
    j = length(mr.m)
    
    ...

    if j == 0 
          return nothing
    end

    ...

    return a, a # state equals the last iterate
end
Base.IteratorSize(::MixedRadix) = Base.SizeUnkown()
3 Likes
3

Thanks for the reply…it raises another question for me. Why do I need a struct? It’s always going to be an array of Ints…

4

Because you want to tell the compiler it’s not just any array of Ints, but ones with the special meaning. Normally, iterate() defined for arrays gives elements of array in order, the behavior you probably don’t want to change.

1 Like
5

If you are trying to learn the iterator protocol with this example, then listen to the others. But if you want to do this in your own code with high efficiency, I believe mixed-radix iteration is exactly what CartesianIndices is for.

ind = CartesianIndices((4, 2, 3))
display.(ind)

Or to be less concise in an example without broadcasting


for i in eachindex(ind)
    display(i)
end
2 Likes
6

The was really helpful thanks. Indeed, I wasn’t so interested in making an iterator as in generating all n-tuples of a set. CartesianIndices seems to be just what I need.

7

Ranges of the CartesianIndex type indeed already do this for you (using a very efficient implementation involving generated functions). It loops in the opposite order from mixed-radix increments, but you can just call reverse to put it in radix order if needed:

julia> for i in CartesianIndex(0,0,0):CartesianIndex(1,2,1)
          println(reverse(Tuple(i)))
       end
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(0, 2, 0)
(0, 2, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)
(1, 2, 0)
(1, 2, 1)

(I don’t want to discourage you from learning how to implement iterators, but you are unlikely to beat the efficiency of CartesianIndices — it is implemented in pure-Julia code, but by an expert (@tim.holy).)

Note that using this constructor gives you 1-based indices, which may or may not be what you want. My example above gives you 0-based indices identical to your desired mixed-radix iteration.

7 Likes