Given a binary tree, you need to compute the length of the diameter of the tree. The diameter of a binary tree is the length of the longest path between any two nodes in a tree. This path may or may not pass through the root.
Example:
Given a binary tree
1
/ \
2 3
/ \
4 5
Return 3, which is the length of the path [4,2,1,3] or [5,2,1,3].
Note: The length of path between two nodes is represented by the number of edges between them.
The description was taken from https://leetcode.com/problems/diameter-of-binary-tree/.
#O(N) Time, O(N) Space for Unbalanced Tree
class Solution:
def diameterOfBinaryTree(self, root: TreeNode) -> int:
diameter = 0
def postorder(node: TreeNode) -> int:
nonlocal diameter
if node is None:
return 0
left = postorder(node.left)
right = postorder(node.right)
height = max(left, right) + 1
diameter = max(diameter, left + right)
return height
postorder(root)
return diameter
For tree problems, it's best to ask the question, "Whose information do I need first?". This will tell us which kind of traversal we need.
For questions involving height, we would need the leaf nodes' information first because, in order for us to know the height of a node, we'd need to know the height of each child node below it. Height isn't how far a node is from the root, but how many nodes are under them.
So we've got our traversal picked, a postorder traversal, what are we trying to figure out? The maximum diameter we'd be able to calculate is the longest two child heights of any given node. This may be two child nodes of a subtree, or from the root, that is why the problem description states it may travel through the root or not.
We can solve this by performing a postorder Depth-First Search on the tree and within each function call, we will calculate the maximum height of the current node and also calculate the maximum sum for every two child heights.
Let's start by initializing a running diameter sum.
diameter = 0
Next, we'll make our postorder function which will take a node as an argument and will return a height integer.
def postorder(node: TreeNode) -> int:
We'll need to clarify that this diameter integer we're referencing is outside the scope of the current function for Python.
We could just include this within the function's return value, but we'd have to return a list instead of an integer so it's cleaner this way.
nonlocal diameter
Let's set our base case.
We're going to be traveling deep down each child leaf of a branch first before coming back up. Once we hit a null node, we will return a height of zero. This way, each parent node can know what their child nodes' height is to know what height they are.
if node is None:
return 0
Otherwise, we will continue our postorder DFS on the left and right child nodes for any given node.
left = postorder(node.left)
right = postorder(node.right)
We'll recurse the left and right branches of the tree, and once the DFS yoyo has come back up so that we can find out what their heights are, we will add one to the max height between them. That will be the height of the current function call's node.
We take the max between them because if we had two children on one branch and one child on the other, we would be as tall as the two children, plus ourselves.
height = max(left, right) + 1
Once we know the child's heights, we will also calculate the current diameter of this node by adding those heights together. When we have that, we can compare that to the maximum diameter we have saved, updating it to the current one if it is greater.
diameter = max(diameter, left + right)
When we are done with those calculations, we will return the height of the current node.
return height
Now that we have our postorder function built, we just need to make the initial call on the root and return the maximum diameter we found.
postorder(root)
return diameter