pythontreedepth-first-searchbreadth-first-search

Yield depth on a depth-first non-recursive tree walker?


I'm trying to find a non-recursive "powerful/versatile" tree walker algorithm, ultimately yielding not just the node but the depth of the node, its parent and its sibling index, and able to use a breadth-first search (BFS) or depth-first search (DFS).

It is possible to combine depth-first and bread-first like this, just yielding the node (NB assumes a node which may or may not have the key CHILD_NODES). Example using Python - "Python" tag added but obviously not specific:

def get_walker(walk_by_depth=True):
    def walk(root):
        queue = [root]
        while len(queue) > 0:
            # the first item in the queue
            node_to_yield = queue.pop(0)

            yield node_to_yield

            if CHILD_NODES in node_to_yield:
                depth += 1    
                new_nodes = node_to_yield[CHILD_NODES]
                if walk_by_depth:
                    queue = new_nodes + queue
                else:
                    queue.extend(new_nodes)
    return walk

... and adding a small amount of code lets you yield a depth for a breadth-first search only:

def get_walker(walk_by_depth=True):
    def walk(root):
        queue = [root]
        depth = 0 
        depth_map_of_remaining_items = {0: 1}
        while len(queue) > 0:
            node_to_yield = queue.pop(0)
            if depth_map_of_remaining_items[depth] == 0:
                depth += 1
            depth_map_of_remaining_items[depth] -= 1

            yield node_to_yield, depth

            if CHILD_NODES in node_to_yield:
                depth += 1    
                new_nodes = node_to_yield[CHILD_NODES]
                if walk_by_depth:
                    queue = new_nodes + queue
                else:
                    queue.extend(new_nodes)
                    if depth not in depth_map_of_remaining_items:
                        depth_map_of_remaining_items[depth] = 0
                    depth_map_of_remaining_items[depth] += len(new_nodes)    
                depth -= 1
    return walk

The above code doesn't in fact work with walk_by_depth=True: it doesn't raise an Exception, it just gets the depth wrong. Instinctively I've got a feeling that the code probably just needs a minimal tweak to yield (correct) depth on a DFS, but I've had no success so far.

The thing is, if I can find a way of yielding the depth with a DFS, I believe it will be a fairly easy step to yield, for both BFS and DFS, the parent node, by maintaining a "depth-to-last-node" map. If you can get the parent you can also get the sibling index, at the simplest by using list's index method (though there may be far cleverer ways to get both the parent and the sibling index...).


Solution

  • The idea is to place the depth in the queue too: put pairs of nodes with their depths in the queue, and it will be easy.

    Without changing more than necessary in your code to make it work, we get this:

    def get_walker(walk_by_depth=True):
        def walk(root):
            queue = [(root, 0)]  # pair the root with its depth
            while len(queue) > 0:
                # the first item in the queue
                node_to_yield, depth = queue.pop(0)  # pop the pair
    
                yield node_to_yield, depth
    
                if CHILD_NODES in node_to_yield:
                    # create pairs for all child nodes
                    new_nodes = [(child, depth+1) for child in node_to_yield[CHILD_NODES]]
                    if walk_by_depth:
                        queue = new_nodes + queue
                    else:
                        queue.extend(new_nodes)
        return walk
    

    Note that pop(0) and new_nodes + queue are not efficient operations. Consider using a deque.