transform_bst_sum_tree.py

from __future__ import annotations


class Node:
    """
    prints the inorder Traversal of transformed tree
    >>> sum = 0
    >>> root = Node(11)
    >>> root.left = Node(2)
    >>> root.right = Node(29)
    >>> root.left.left = Node(1)
    >>> root.left.right = Node(7)
    >>> root.right.left = Node(15)
    >>> root.right.right = Node(40)
    >>> root.right.right.left = Node(35)
    >>> printInorder(root)
    1 2 7 11 15 29 35 40

    >>> transformTree(root)

    >>> printInorder(root)
    139 137 130 119 104 75 40 0
    """

    def __init__(self, number: int) -> None:
        self.data = number
        self.left = None
        self.right = None


# Recursive function to transform a BST to sum tree.
# This function traverses the tree in reverse inorder so
# that we have visited all greater key nodes of the currently
# visited node
def transform_tree_util(root: Node | None) -> None:
    """
    Transform a binary tree into a sum tree.

    Example:
    >>> root = Node(11)
    >>> root.left = Node(2)
    >>> root.right = Node(29)
    >>> transformTree(root)
    >>> root.data
    60
    >>> root.left.data
    31
    >>> root.right.data
    29
    """

    # Base case
    if root == None:
        return

    # Recur for right subtree
    transform_tree_util(root.right)

    # Update sum
    global sum
    sum = sum + root.data

    # Store old sum in the current node
    root.data = sum - root.data

    # Recur for left subtree
    transform_tree_util(root.left)


# A wrapper over transformTreeUtil()
def transform_tree(root: Node | None) -> None:
    """
    Transform a binary tree into a sum tree.

    Example:
    >>> root = Node(11)
    >>> root.left = Node(2)
    >>> root.right = Node(29)
    >>> transformTree(root)
    >>> root.data
    60
    >>> root.left.data
    31
    >>> root.right.data
    29
    """

    # sum = 0 #Initialize sum
    transform_tree_util(root)


# A utility function to prindorder traversal of a
# binary tree
def print_inorder(root: Node | None):
    """
    Perform an inorder traversal of a binary tree and print the nodes.

    Example:
    >>> root = Node(11)
    >>> root.left = Node(2)
    >>> root.right = Node(29)
    >>> printInorder(root)
    2 11 29
    """

    if root == None:
        return

    print_inorder(root.left)
    print(root.data, end=" ")
    print_inorder(root.right)


# Driver Program to test above functions
if __name__ == "__main__":
    sum = 0
    root = Node(11)
    root.left = Node(2)
    root.right = Node(29)
    root.left.left = Node(1)
    root.left.right = Node(7)
    root.right.left = Node(15)
    root.right.right = Node(40)
    root.right.right.left = Node(35)

    print_inorder(root)

    transform_tree_util(root)
    print_inorder(root)