Source code for example_code.graphs

"""A simple tree implementation with basic pre- and post-visitors."""


class TreeNode:
    """A basic tree implementation.

    A tree is simply a collection of connected TreeNodes.

    Parameters
    ----------
    value:
        An arbitrary value associated with this node.
    children:
        The TreeNodes which are the children of this node.
    """

    def __init__(self, value, *children):
        self.value = value
        self.children = children

    def __repr__(self):
        """Return the canonical string representation."""
        return f"{type(self).__name__}{(self.value,) + self.children}"

    def __str__(self):
        """Serialise the tree recursively as parent -> (children)."""
        childstring = ", ".join(map(str, self.children))
        return f"{self.value!s} -> ({childstring})"


def previsitor(tree, fn, fn_parent=None):
    """Traverse tree in preorder applying a function to every node.

    Parameters
    ----------
    tree: TreeNode
        The tree to be visited.
    fn: function(node, fn_parent)
        A function to be applied at each node. The function should take
        the node to be visited as its first argument, and the result of
        visiting its parent as the second.
    """
    fn_out = fn(tree, fn_parent)

    for child in tree.children:
        previsitor(child, fn, fn_out)


def postvisitor(tree, fn):
    """Traverse tree in postorder applying a function to every node.

    Parameters
    ----------
    tree: TreeNode
        The tree to be visited.
    fn: function(node, *fn_children)
        A function to be applied at each node. The function should take the
        node to be visited as its first argument, and the results of
        visiting its children as any further arguments.
    """
    return fn(tree, *(postvisitor(c, fn) for c in tree.children))