Tree traversal without recursion

#algorithms #performance

So, you've got a pretty deep hierarchy, or maybe want to have a better debugging experience without seeing all the noise created by your recursive function on the call stack. You decide to convert it (such as one below) to using an explicit stack.

interface TreeNode {
    name: string;
    children: TreeNode[];
}

function walkTreeRecursive(node: TreeNode): void {
    console.log("begin", node.name); // Pre-order
    for (const child of node.children) {
        walkTreeRecursive(child);
        console.log("yield", node.name); // In-order
    }
    console.log("end", node.name); // Post-order
}

Traversal

Let's start with the naive approach, the one you've already thought of, or are likely to see on stackoverflow. The obvious downside of it is that we can't run any in-order or post-order code. For example, we can no longer increment the index of the next child element and fetch it from the array on demand. Instead, we have to queue them all at once.

function walkTreePreNaive(rootNode: TreeNode): void {
    const stack: TreeNode[] = [rootNode];

    while (stack.length > 0) {
        const node = stack.pop()!;

        console.log("begin", node.name); // Pre-order

        // NAIVE adds all children at once,
        // therefore uses more memory than the recursive approach.
        // Siblings from parent/current levels will sit on the stack
        // and wait their turn for quite a while.
        for (const child of node.children.reverse()) {
            stack.push(child);
        }
    }
}

If we started using a stack to mimic the function call stack, we should also know about stack frames. There, local variables can be saved before switching to process any child. By introducing an index variable, we will regain all the missing functionality.

class StackFrame {
    constructor(
        public node: TreeNode,
        public childIndex: number = 0,
    ) {}
}

function walkTreeWithIndex(rootNode: TreeNode): void {
    const stack: StackFrame[] = [new StackFrame(rootNode)];

    while (stack.length > 0) {
        const frame = stack.pop()!;
        const { node } = frame;

        if (frame.childIndex === 0) {
            console.log("begin", node.name); // Pre-order
        } else {
            console.log("yield", node.name); // In-order
        }

        if (frame.childIndex < node.children.length) {
            // add self for in-order
            frame.childIndex += 1;
            stack.push(frame);
            // add child for pre-order
            stack.push(new StackFrame(node.children[frame.childIndex - 1]));
        } else {
            console.log("end", node.name); // Post-order
        }
    }
}

That's it. It's more complicated, and less readable than using recursion, but it gets the job done.

Bonus: aggregate/reduce/transform

What if a recursive function needs a return value for the result of the aggregation? No problem, just add another "local variable" to the stack-frame and a place to store the return type so the caller can access it.

 class StackFrame {
     constructor(
         public node: TreeNode,
         public childIndex: number = 0,
+        public total: number = 0,
     ) {}
 }

 function transformTreeStack(rootNode: TreeNode): void {
     const stack: StackFrame[] = [new StackFrame(rootNode)];

+    let retValue!: number;
     while (stack.length > 0) {
         const frame = stack.pop()!;
         const { node } = frame;

         if (frame.childIndex === 0) {
             console.log("begin", node.name); // Pre-order
         } else {
+            frame.total += retValue;
         }

         if (frame.childIndex < node.children.length) {
             // add self for in-order
             frame.childIndex += 1;
             stack.push(frame);
             // add child for pre-order
             stack.push(new StackFrame(node.children[frame.childIndex - 1]));
         } else {
+            retValue = frame.total + node.name.length;
         }
     }
+    console.log("total", retValue);
 }