Recursive and Level-Order Traversal for N-ary Trees

Once you finish this piece, you will have internalized the core patterns and can immediately apply them to the problems listed below:

Prerequisites

Make sure you are comfortable with the following before continuing:

Recursive and level-order traversal of binary trees

Summary in One Sentence

An N-ary (multi-way) tree is simply the binary-tree idea extended to an arbitrary number of children; a binary tree is therefore just the special case of an N-ary tree that is limited to two children. Traversal of an N-ary tree follows the same logic, merely adapted to that more general structure. A forest is a collection of one or more N-ary trees; a lone N-ary tree can be regarded as a degenerate forest.

A binary-tree node is defined with at most two children:

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

An N-ary-tree node stores a value together with a list of any number of children:

class Node:
    def __init__(self, val: int):
        self.val = val
        self.children = []

That single change in the node layout is the only structural distinction that matters.

Forest

The term “forest” is worth a short explanation because it reappears when we later examine the Union-Find (Disjoint-Set) algorithm.

A forest is nothing more than a set of N-ary trees; a solitary tree counts as a special case of a forest. In code it is simply a list of the root nodes of those trees:

List<Node> forest;

To visit every node that belongs to the forest you start an independent DFS or BFS from each root.

Inside a Union-Find structure we simultaneously keep the roots of many N-ary trees; the collection of those roots is precisely a forest.

We now turn to the two classic ways of walking an N-ary tree—exactly the same pair of techniques used for binary trees: recursive (DFS) traversal and level-order (BFS) traversal.

Recursive Traversal (DFS)

Place the familiar binary-tree skeleton next to the N-ary version:

# Binary-tree traversal skeleton
def traverse_binary_tree(root):
    if root is None:
        return
    # pre-order position
    traverse_binary_tree(root.left)
    # in-order position
    traverse_binary_tree(root.right)
    # post-order position


# N-ary-tree traversal skeleton
def traverse_n_ary_tree(root):
    if root is None:
        return
    # pre-order position
    for child in root.children:
        traverse_n_ary_tree(child)
    # post-order position

The only real difference is that the N-ary skeleton has no “in-order” slot. Because a node may possess an arbitrary number of children, the idea of visiting the node “between” its children ceases to have a clear meaning.

The interactive visualization that follows lets you click the indicated line repeatedly and watch the visit order; you will see that it closely mirrors the order produced by a binary-tree walk.

Level-Order Traversal (BFS)

Level-order traversal of an N-ary tree is performed with a queue in precisely the same way as for a binary tree; the only mechanical change is that every child of the current node is enqueued instead of merely the left and right children. Three common implementations exist; the first two are shown below.

First Implementation

This version does not record depth:

from collections import deque

def level_order_traverse(root):
    if root is None:
        return
    
    q = deque()
    q.append(root)
    
    while q:
        cur = q.popleft()
        # visit the current node
        print(cur.val)
        
        # enqueue every child of the current node
        for child in cur.children:
            q.append(child)

Clicking the highlighted line in the visualization panel several times reveals the step-by-step progress of the level-order walk.

Second Implementation

This version also tracks the depth of each node (the root is treated as depth 1):

from collections import deque

def levelOrderTraverse(root):
    if root is None:
        return
    q = deque()
    q.append(root)
    # depth of the level currently being processed (root = 1)
    depth = 1

    while q:
        sz = len(q)
        for i in range(sz):
            cur = q.popleft()
            # visit the node while knowing its depth
            print(f"depth = {depth}, val = {cur.val}")

            for child in cur.children:
                q.append(child)
        depth += 1

Solution of 429. N-ary Tree Level Order Traversal: Given an n-ary tree, return the level order traversal of its nodes' values.

class Solution:
    def levelOrder(self, root: 'Node') -> List[List[int]]:
        res = []

        # base case
        if not root:
            return res

        q = [root]

        while len(q)!=0:
            size = len(q)
            layer = []
            for i in range(size):
                node = q[i]
                layer.append(node.val)
                for child in node.children:
                    if child is not None:
                        q.append(child)
            q = q[size:]
            res.append(layer)
        
        return res