pythontreetree-traversalmptt

Stack based modified preorder tree traversal


I have a recursive implementation of a modified preorder tree traversal (nested set model) that I want to implement without the use of recursion.

from collections import deque

def mptt_recurse(tree, node, preorder=None):

    if node not in tree: return
    if preorder is None: preorder = deque()

    for child in tree[node]:
        preorder.append(child)
        mptt_recurse(tree, child, preorder)
        preorder.append(child)

    return preorder

The recursive implementation works as expected:

>>> tree = {
    "root": ["food"],
    "food": ["meat", "veg"],
    "meat": ["pork", "lamb"],
    "pork": ["bacon", "sausage"],
    "lamb": ["cutlet"],
    "soup": ["leak", "tomato"]
} 
>>> mptt_recuser(tree, "root")
deque(['food', 'meat', 'pork', 'bacon', 'bacon', 'sausage', 'sausage', 'pork', 'lamb', 'cutlet', 'cutlet', 'lamb', 'meat', 'veg', 'veg', 'food'])   

Each node appears twice with left and right values represented by the position in the deque.

def mptt_stack(tree, node):

    if node not in tree: return
    stack = deque(tree[node])
    preorder = deque()

    while stack:
        node = stack.pop()
        preorder.append(node)
        if node not in tree:
            continue
        for child in reversed(tree[node]):
            stack.append(child)

    return preorder

With my stacked based implementation I have only been able to figure out how to get the preorder (not the modified preorder with both the left and right labels for each node).

>>> mptt_stack(tree, "root")
deque(['food', 'meat', 'pork', 'bacon', 'sausage', 'lamb', 'cutlet', 'veg'])

Any pointers on a non recursive implementation would be great.


Solution

  • This will give the same results as mptt_recurse. It keeps an additional piece of information on the stack, which indicates whether it's entering or leaving the node:

    def mptt_stack(tree, node):
        if node not in tree: return
        preorder = deque()
    
        stack = []
        for child in reversed(tree[node]):
            stack.append([child, True])
    
        while stack:
            (node, first) = stack.pop()
            preorder.append(node)
            if first:
                stack.append([node, False])
                if node in tree:
                    for child in reversed(tree[node]):
                        stack.append([child, True])
    
        return preorder
    

    If you want to include the initial node in the result, you can replace the loop that initializes stack with:

    stack = [[node, True]]