Traversing Trees with Generators

Introduction

Earlier this year, I was interviewing for frontend engineering positions and tree questions were very popular. After repetitively solving tree recursion problems and trying to improve my tactics, I came across a novel JavaScript feature – Generators – that makes tree recursion as easy as array iteration.

What Are Generators?

Why do we even need Generators? To put it bluntly, implementing the Iteration protocols for an object or a class is a PITA. With generators, we get a concise syntax that makes it easy to implement the Iteration protocols – which we’ll leverage later on to simplify tree recursion into iteration.

A complete explanation of Generators is beyond the scope of this article but it’s important to understand a few basics of the syntax to understand the following exercises.

  1. A Generator is a special type of function that can yield multiple values whereas a regular function can only return a single value.
  2. The function* syntax allows us to define a function that will return a Generator, and the yield syntax allows us to return the next value from inside of a Generator.
  3. Generators implement the Iteration protocols so we can leverage for...of loops.

For example, here’s a function that returns a Generator for and infinite series of numbers. Notice that this function will never terminate but by wrapping it in a Generator we can pause and restart it whenever we want.

1
2
3
4
5
6
7
8
9
function* numbers() {
  let n = 1;
  while (true) yield n++
}

const gen = numbers();

console.log(gen.next());
console.log(gen.next());
>> Object { value: 1, done: false }
>> Object { value: 2, done: false }

Exercises

For all of the exercises below, we’ll be using the tree above to illustrate each type of traversal. The tree above was created by using a helper function called createTreeNode which accepts a single value and an array of children.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
function createTreeNode(value, children = []) {
  const node = {
    value,
    children
  };
  // Create references to the parent node.
  for (let child of node.children) {
    child.parent = node;
  }
  return node;
};

const tree = createTreeNode(1, [
  createTreeNode(2, [
    createTreeNode(3),
    createTreeNode(4),
  ]),
  createTreeNode(5, [
    createTreeNode(6),
  ]),
]);
Depth-first Traversal

Every single tree question that I was given involved depth-first traversal. For example, flatten an array of trees, search the DOM for elements with a specific class name, or build a recursive React component.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
function* traverseTree(node) {
  yield node;

  for (let child of node.children) {
    yield* traverseTree(child);
  }
}

for (let node of traverseTree(tree)) {
  console.log(node.value);
}
1
2
3
4
5
6

Breadth-first Traversal

I was never asked any questions that required breadth-first traversal but I’m including this exercise because it shows the versatility of this technique. For example, print the values of a tree in row order.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
function* traverseTreeBreadthFirst(node) {
  let queue = [node]

  while (queue.length > 0) {
    let curr = queue.shift()

    for (let child of curr.children) {
      queue.push(child);
    }

    yield curr;
  }
}

for (let node of traverseTreeBreadthFirst(leaf)) {
  console.log(node.value);
}
1
2
5
3
4
6

Backtracking

I encountered backtracking once as an extension to a depth-first tree traversal question. For example, search the DOM for leaf elements with a specific hierarchy of class names.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
function* backtrackTree(node) {
  let curr = node

  while (curr) {
    yield curr

    curr = curr.parent
  }
}

// Find the rightmost tree leaf node.
const leafNode = [...traverseTree(tree)].pop();

for (let node of backtrackTree(leafNode)) {
  console.log(node.value);
}
6
5
1

Further Exercises
  • Can you flatten an array of trees?
  • Can you implement map, filter and reduce for trees?
  • Can you implement iteration for another data structure?

Conclusion

Hopefully, I’ve made a strong case that tree recursion can be as easy as array iteration by using JavaScript Generators. All of the source code for this blog post is available on GitHub.

Thank you to Michael Geraci for editing.