in Lua, how to write an iterator for a binary tree?

Question

In Lua, how to write an iterator for a binary tree. e.g. so I could do something like:

for node in tree:visit() do ... end

I've been trying but best I can do is to walks the tree, carrying around some function fun as a payload. e.g.

function klass(s,    t) t={ako=s}; t.__index=t; return t end

NODE=klass"NODE"
function NODE.new(x,l,r) return setmetatable( {x=x, left=l, right=r},NODE) end

function NODE:visit(fun,  lvl)
  lvl=lvl or 0
  fun(lvl,self)
  for _,kid in pairs{self.left, self.right} do 
    if kid then kid:visit(fun, lvl+1) end end end

--- eg -----------------------------
function eg(x,stop,     node)
  node = NODE.new(x)
  if x < stop/2 then
    node.left = eg(2*x,stop)
    node.right = eg(2*x+1,stop) end
  return node end 

tree=eg(1,32)
tree:visit( function(lvl,node) print(('|.. '):rep(lvl)..node.x) end )

This produces the output I want:

1
|.. 2
|.. |.. 4
|.. |.. |.. 8
|.. |.. |.. |.. 16
|.. |.. |.. |.. 17
|.. |.. |.. 9
|.. |.. |.. |.. 18
|.. |.. |.. |.. 19
|.. |.. 5
|.. |.. |.. 10
|.. |.. |.. |.. 20
|.. |.. |.. |.. 21
|.. |.. |.. 11
|.. |.. |.. |.. 22
|.. |.. |.. |.. 23
|.. 3
|.. |.. 6
|.. |.. |.. 12
|.. |.. |.. |.. 24
|.. |.. |.. |.. 25
|.. |.. |.. 13
|.. |.. |.. |.. 26
|.. |.. |.. |.. 27
|.. |.. 7
|.. |.. |.. 14
|.. |.. |.. |.. 28
|.. |.. |.. |.. 29
|.. |.. |.. 15
|.. |.. |.. |.. 30
|.. |.. |.. |.. 31

But how do I write this such that I can work it into a for loop?

There are two solutions I'd like to avoid:

• I don't want to use pairs to over a list of nodes generated via the my code above (since I'd like to support incremental traversal of trees of arbitrary depth)
• I want the answer to be (much) less complex than https://github.com/guicaulada/lua-bintrees/blob/main/bintrees/iterator.lua

Answer 1

If you need to convert a recursive function into a loop, then you need a stack to simulate the call stack during recursion.

The for-in statement in lua needs 3 values: an iterator function, a state (the stack here), an initial value (omitted here):

function NODE:visit()
  return function (stack)
    if #stack == 0 then return nil end
    local node = table.remove(stack)
    if node.right then stack[#stack + 1] = node.right end
    if node.left then stack[#stack + 1] = node.left end
    return node;
  end, {self}
end

And the usage:

for node in tree:visit() do
  print(node.x)
end

Answer 2

Avoiding recursion and writing an explicitly iterative iterator is one option. If the depth of your binary tree is reasonably bounded, your recursive implementation works just as well. If you want an actual iterator, you can easily obtain one by passing coroutine.yield as fun and "wrapping" the iteration in a coroutine via coroutine.wrap:

for lvl, node in coroutine.wrap(function() tree:visit(coroutine.yield) end) do
    print(('|.. '):rep(lvl)..node.x)
end

Of course you could also write a "method" for this:

function NODE:nodes()
     local function visit(node)
         coroutine.yield(node)
         if node.left then visit(node.left) end
         if node.right then visit(node.right) end
     end
     return coroutine.wrap(function() visit(self) end)
end
-- Usage:
for node in tree:nodes() do
    print(node.x)              
end

A recursive solution has the drawback of being limited by stack size (though this limit is very large on newer Lua versions, and can again be circumvented using "trampolines" or even more coroutines), and coroutines add some overhead both time- and space-wise.

Most of the time this isn't critical though, and I find that coroutines yield very elegant solutions for obtaining an "iterative" iterator from a "recursive" one.

In this example, the difference isn't very large, because you never need to go "back" to a node (you're doing a "preorder" traversal: You visit the root first, then visit the children).

As an exercise, try to implement an inorder traversal. In the "recursive" / coroutine-based version, this should be trivial. In the iterative version, it's not quite as easy.