Practice: An Iterator Over a Binary Tree

Walk the tree in order without exposing its nodes, then loop over it twice.

A binary tree is a collection with no index, which makes it a natural place for an Iterator. This exercise asks you to write one as a generator and make it reusable. Work out your answer before opening the solution.

A Node holds a value and an optional left and right node. A BinaryTree holds a root node. Make BinaryTree iterable so that:

for value in tree returns the values in order, smallest first. In a binary search tree, that means the left subtree, then the node, then the right subtree.
The caller never sees a Node.
list(tree) gives the same result the second time it is called.
Built-in functions such as max work on the tree.

Use this tree: 5 at the root; 3 on its left, with 1 and 4 below it; 8 on its right, with 9 to its right.

Show a solution

from collections.abc import Iterator
from dataclasses import dataclass


@dataclass
class Node:
    value: int
    left: "Node | None" = None
    right: "Node | None" = None


class BinaryTree:
    def __init__(self, root: Node | None) -> None:
        self._root = root

    def __iter__(self) -> Iterator[int]:
        yield from self._in_order(self._root)

    def _in_order(self, node: Node | None) -> Iterator[int]:
        if node is None:
            return
        yield from self._in_order(node.left)
        yield node.value
        yield from self._in_order(node.right)


tree = BinaryTree(Node(5, Node(3, Node(1), Node(4)), Node(8, None, Node(9))))
print("first: ", list(tree))
print("second:", list(tree))
print("largest:", max(tree))

OutputRan with python:3.12-slim (Python 3.12.13)

first:  [1, 3, 4, 5, 8, 9]
second: [1, 3, 4, 5, 8, 9]
largest: 9

`__iter__` is a generator, so every walk starts at the root.

__iter__ is a generator method, so each list(tree) starts a fresh walk, and the second pass gives the same six values. The recursive helper uses yield from, which passes on every value of the inner generator, to visit the left subtree, the node and the right subtree in turn. The caller receives integers only, never a Node, and max works because it accepts any iterable.

An iterator class would need a stack of nodes kept by hand, to remember where to return after each left subtree. The generator keeps that stack for you, in its paused recursive calls.

← → move through sections and cards