Welcome, guest | Sign In | My Account | Store | Cart

Simple breadth-first, depth-first tree traversal (Python recipe) by Jack Trainor
ActiveState Code (http://code.activestate.com/recipes/579138/)

When you want to avoid recursion with a tree, you read the tree nodes into a stack, which is organized either breadth-first or depth-first.

Here are two dead simple routines for doing so. Most of the recipe is just a test bed for those functions.

Python, 96 lines
Download
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
"""
recipe_simple_tree_traversal.py

Simple breadth-first and depth-first tree traversal for any tree.

The recipe is contained within the first two functions. The rest is
a test bed for demonstration.
"""
__author__ = "Jack Trainor"
__date__ = "2015-12-16"

import sys
    
def get_breadth_first_nodes(root):
    nodes = []
    stack = [root]
    while stack:
        cur_node = stack[0]
        stack = stack[1:]
        nodes.append(cur_node)
        for child in cur_node.get_children():
            stack.append(child)
    return nodes

def get_depth_first_nodes(root):
    nodes = []
    stack = [root]
    while stack:
        cur_node = stack[0]
        stack = stack[1:]
        nodes.append(cur_node)        
        for child in cur_node.get_rev_children():
            stack.insert(0, child)
    return nodes

########################################################################
class Node(object):
    def __init__(self, id_):
        self.id = id_
        self.children = []
        
    def __repr__(self):
        return "Node: [%s]" % self.id
    
    def add_child(self, node):
        self.children.append(node) 
    
    def get_children(self):
        return self.children         
    
    def get_rev_children(self):
        children = self.children[:]
        children.reverse()
        return children         

########################################################################
def println(text):
    sys.stdout.write(text + "\n")
    
def make_test_tree():
    a0 = Node("a0")
    b0 = Node("b0")      
    b1 = Node("b1")      
    b2 = Node("b2")      
    c0 = Node("c0")      
    c1 = Node("c1")  
    d0 = Node("d0")   
    
    a0.add_child(b0) 
    a0.add_child(b1) 
    a0.add_child(b2)
    
    b0.add_child(c0) 
    b0.add_child(c1) 
    
    c0.add_child(d0)
    
    return a0                  

def test_breadth_first_nodes():
    root = make_test_tree()
    node_list = get_breadth_first_nodes(root)
    for node in node_list:
        println(str(node))

def test_depth_first_nodes():
    root = make_test_tree()
    node_list = get_depth_first_nodes(root)
    for node in node_list:
        println(str(node))

########################################################################
if __name__ == "__main__":
    test_breadth_first_nodes()
    println("")
    test_depth_first_nodes()

Recursion is natural for reading trees, but I like to eliminate recursion when I can. Fortunately there is a standard CompSci solution which is to read the tree into a node stack organized breadth-first or depth-first.

Unfortunately most of the online code examples are written in Lisp or using advanced Python features which obscure what is really going on.

This recipe could be rewritten with deques or iterators or generators or what-have-you, but I kept it to the most basic Python features so even novices can see what's happening, which frankly isn't all that much.

Enjoy!

Tags: tree