Develop a Custom Iterator for Tree data structures in Python

Write a Python program to develop a custom iterator that iterates over a tree data structure.

The task is to create a Python program that develops a custom iterator for traversing a tree data structure. This iterator will systematically visit each node in the tree, enabling easy and efficient access to the tree's elements in a specific order (e.g., pre-order, in-order, or post-order traversal). The program will define the tree structure and the custom iterator class, providing a flexible and reusable way to iterate over trees in various applications.

Sample Solution:

Python Code :

# Define a class representing a node in a tree
class TreeNode:
    def __init__(self, value):
        self.value = value
        self.children = []  # Initialize an empty list to store child nodes

    def add_child(self, child_node):
        self.children.append(child_node)  # Add a child node to the list of children

# Define a class implementing a custom iterator for tree traversal
class TreeIterator:
    def __init__(self, root):
        self.stack = [root]  # Initialize stack with root node

    def __iter__(self):
        return self

    def __next__(self):
        if not self.stack:  # If stack is empty, no more nodes to traverse
            raise StopIteration
        node = self.stack.pop()  # Get the top node from stack
        for child in reversed(node.children):  # Add children to stack in reverse order
            self.stack.append(child)
        return node.value

# Example usage:
if __name__ == "__main__":
    # Create a tree with a root node and some child nodes
    root = TreeNode(1)
    root.add_child(TreeNode(2))
    root.add_child(TreeNode(3))
    root.children[0].add_child(TreeNode(4))
    root.children[0].add_child(TreeNode(5))
    root.children[1].add_child(TreeNode(6))

    # Create a TreeIterator instance and iterate over the tree
    tree_iterator = TreeIterator(root)
    for value in tree_iterator:
        print(value)

Output:

1
2
4
5
3
6

Explanation:

• The code defines two classes: "TreeNode" and "TreeIterator".
• TreeNode: Represents a node in a tree structure. Each node has a value and can have child nodes. The "add_child()" method adds a child node to the current node.
• TreeIterator: Implements a custom iterator for traversing a tree in a depth-first manner. It uses a stack to keep track of nodes to visit. The "init()" method initializes the iterator with the root node. The "iter()" method returns the iterator object itself. The "next()" method implements logic for traversing the tree. It pops a node from the stack, visits its children in reverse order, and returns the value of the node.
• Finally the example usage section demonstrates how to create a tree, populate it with nodes, create a 'TreeIterator' instance with the root node, and iterate over the tree using the custom iterator, printing the values of the nodes in the order they are visited.

Python Code Editor :

Have another way to solve this solution? Contribute your code (and comments) through Disqus.

Previous: Create a Python Class-based Decorator to Log method execution time.
Next: Matrix Multiplication in Python Using list comprehensions