Search code examples
pythonbinary-tree

How to update height of all nodes in a binary tree with one function call?


I have a binary tree consisting of nodes of this class:

class Node:
    def __init__(self,key,parent=None,right=None,left=None):
        self.key = key
        self.left = left
        self.right = right
        self.parent = parent
        self.height = 0

An example of a tree construction could be:

a = Node(1)
a.left = Node(2,parent=a)
a.right = Node(5,parent=a)
a.right.right = Node(3,parent=a.right)

I want to write a function that updates the heights of all nodes in this tree. I have found ways to get the height of one node. However, I want to update all heights at once.

How could I do this?


Solution

  • You could add this recursive method:

    def syncheight(self):
        self.height = max(
            -1 if self.left is None else self.left.syncheight(), 
            -1 if self.right is None else self.right.syncheight()
        ) + 1
        return self.height
    

    You would call it on the root as follows:

    a.syncheight()
    

    Note that I would call the class Node instead of node, and you have a different parameter name than what you use in the body of the constructor (left, leftChild)

    Suggestion

    I would not pass the parent to the constructor, as that should follow from the information passed as left and right. The latter should get their parent property updated to refer to the currently created node.

    You can then also keep height updated as you go, by updating the heights of the ancestors of the newly created node:

    class Node:
        def __init__(self, key, left=None, right=None):
            self.key = key
            self.left = left
            self.right = right
            self.parent = None
            if left:
                left.parent = self
            if right:
                right.parent = self
            self.height = -1
            self.bubbleheight()
        
        def bubbleheight(self):
            height = max(
                -1 if self.left is None else self.left.height,
                -1 if self.right is None else self.right.height
            ) + 1
            if height != self.height:
                self.height = height
                if self.parent is not None:
                    self.parent.bubbleheight()
    
    tree = Node(1,
                Node(2),
                Node(5,
                    Node(3)
                )
            )
    
    print(tree.height)  # 2