CODEWITHSUNDEEP
pythontree
Apollo
Apollo

Reputation: 9064

Optimize finding diameter of binary tree in Python

I'm wondering how I can optimally find the diameter (or longest path between any two leaf nodes) of a binary tree. I have the basic solution below, but the second solution requires passing pointers. How can I do something like this in Python?

def find_tree_diameter(node):
    if node == None:
        return 0

    lheight = height(node.left)
    rheight = height(node.right)

    ldiameter = find_tree_diameter(node.left)
    rdiameter = find_tree_diameter(node.right)

    return max(lheight+rheight+1, ldiameter, rdiameter)

def find_tree_diameter_optimized(node, height):
    lheight, rheight, ldiameter, rdiameter = 0, 0, 0, 0

    if node == None:
#        *height = 0;
        return 0

    ldiameter = diameterOpt(root.left, &lheight)
    rdiameter = diameterOpt(root.right, &rheight)

#    *height = max(lheight, rheight) + 1;

    return max(lh + rh + 1, max(ldiameter, rdiameter));

Upvotes: 2

Views: 2273

Answers (2)

Shyambeer Singh
Shyambeer Singh

Reputation: 328

Simple Python 3 solution

def findDepth(root):

    if root is None:
        return 0

    return 1 + max(findDepth(root.left), findDepth(root.right))

class Solution:
    def diameterOfBinaryTree(self, root: TreeNode) -> int:

        if root is None:
            return 0

        left = findDepth(root.left)
        right = findDepth(root.right)

        ldia = self.diameterOfBinaryTree(root.left)
        rdia = self.diameterOfBinaryTree(root.right)


        return max(left+right, max(ldia, rdia))

Upvotes: 1

Paul Hankin
Paul Hankin

Reputation: 58319

Python supports multiple return values, so you don't need pointer arguments like in C or C++. Here's a translation of the code:

def diameter_height(node):
    if node is None:
        return 0, 0
    ld, lh = diameter_height(node.left)
    rd, rh = diameter_height(node.right)
    return max(lh + rh + 1, ld, rd), 1 + max(lh, rh)

def find_tree_diameter(node):
    d, _ = diameter_height(node)
    return d

The function diameter_height returns the diameter and the height of the tree, and find_tree_diameter uses it to just compute the diameter (by discarding the height).

The function is O(n), no matter the shape of the tree. The original function is O(n^2) in the worst case when the tree is very unbalanced because of the repeated height calculations.

Upvotes: 12

Related Questions