Question

I have a recursive method that I’m using to walk over a red black tree, and store various node information (in the list storage).

def _walk (self, storage, func, starting_node) :
    if starting_node is not self._nil :
        self._walk(storage, func, starting_node.left)
        storage.append(func(starting_node))
        self._walk(storage, func, starting_node.right)

However, I’d like to re-implement this method so that it builds a generator (from what I understand this should save both time and memory). What’s the “best” way of doing that?

Answer 1

Begin by decoupling the action from the walking

def _walk (self, starting_node) :
    if starting_node is not self._nil :
        for x in self._walk(starting_node.left):
            yield x
        yield starting_node
        for x in self._walk(starting_node.right):
            yield x

def traverse(self):
    starting_node = ???     # maybe these are passed as
    func = ???              # parameters to traverse
    for node in self._walk(starting_node):
        yield func(node)

traverse is roughly equivalent to

imap(func, self._walk(starting_node))

or this generator expression

(func(x) for x in self._walk(starting_node))

---

You can reduce the amount of stack used by manually optimising the tail recursion

def _walk (self, starting_node) :
    while starting_node is not self._nil :
        for x in self._walk(starting_node.left):
            yield x
        yield starting_node
        starting_node = starting_node.right