#543. Diameter of Binary Tree | LeetCode Python Solution

Problem

LeetCode Problem

543. Diameter of Binary Tree

Python Solution

The idea is that the diameter across a node is the sum of the maximum depths of it’s left subtree and right subtree.

At every node, we compute the height of the left subtree and that of the right subtree. We can then calculate the diameter across this node.

A max_diameter variable is used to keep track of the largest such diameter across all nodes.

Finally, we return max_diameter.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def diameterOfBinaryTree(self, root: Optional[TreeNode]) -> int:
        max_diameter = 0

        # Walk through each tree node and computer its diameter
        stk = [root]
        while len(stk) > 0:
            # Pop tree node from the stack
            c = stk.pop()

            # Current node's diameter = left tree height + right tree height
            curr_node_diameter = self.getTreeHeight(c.left) + self.getTreeHeight(c.right)

            # Update max_diameter if current node's diameter is greater than what we have seen so far
            max_diameter = max(max_diameter, curr_node_diameter)

            # Push left child to the stack
            if c.left is not None:
                stk.append(c.left)

            # Push right child to the stack
            if c.right is not None:
                stk.append(c.right)

        return max_diameter

    def getTreeHeight(self, root: Optional[TreeNode]) -> int:
        # Tree height is 0 when the tree is empty
        if root is None:
            return 0

        # Tree height is 1 when the root node is the only tree node
        if root.left is None and root.right is None:
            return 1

        # Push the root node to the stack and initialize tree height to 1
        tree_height = 1
        node_stack = [[root, tree_height]]
        while len(node_stack) > 0:
            # Pop from the stack
            curr = node_stack.pop()

            # If we just popped a leaf node, find out the tree height
            # Update tree_height variable if it's greater than tree_height we have recorded so far
            if curr[0].left is None and curr[0].right is None:
                tree_height = max(tree_height, curr[1])

            # Push left child along with its level to the stack
            if curr[0].left is not None:
                node_stack.append([curr[0].left, curr[1]+1])

            # Push right child along with its level to the stack
            if curr[0].right is not None:
                node_stack.append([curr[0].right, curr[1]+1])

        return tree_height