leetcode-grind/problems/543/main.py

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def diameterOfBinaryTree(self, root: TreeNode) -> int:
        if root is not None:
            l = self.height(root.left)
            r = self.height(root.right)
            ldiameter = self.diameterOfBinaryTree(root.left)
            rdiameter = self.diameterOfBinaryTree(root.right)
            return max(l + r, max(ldiameter, rdiameter))
        return 0
    
    def height(self, tree):
        if tree is None:
            return 0
        return 1 + max(self.height(tree.left), self.height(tree.right))

s = Solution()
print("Expected: 3")
bt = TreeNode(1)
bt.left = TreeNode(2)
bt.right = TreeNode(3)
bt.left.left = TreeNode(4)
bt.left.right = TreeNode(5)
print("Got:", s.diameterOfBinaryTree(bt))