Tree problems can be a little difficult to wrap your head around, especially when you are a beginner. This post will try and make tree problems easy for you with the help of some examples.
1. Sum of all values
The goal is to get sum of all the values in the tree.
def getSum(root):
if not root:
return 0
return root.val + getSum(root.left) + getSum(root.right)
This is quite simple. The sum of all values in the tree is simply the sum of:
a) the current node
b) the sum of all values on its left
c) the sum of all values on its right
What if there's no node? That should return zero. That's also our base case.
2. Get max value in the tree
The goal here is to get max value from a tree where the values are all >=0.
def getMax(root):
if not root:
return 0
return max(root.val, getMax(root.left), getMax(root.right))
To get the max value in the tree, we return the maximum between:
a) the value of current node
b) the max value on its left
c) the max value on its right
And when there is no node, we return the max as 0. That's also our base case.
3. Get max tree height
The goal here is to find the maximum tree height or the maximum depth of the tree.
For example, the maximum depth of this tree would be 3:
1
/ \
2 3
/ \
4 5
def getHeight(root):
if not root:
return 0
return 1 + max(getHeight(root.left), getHeight(root.right))
To get the max height, we count the current node as 1 and add it to the max height we get from its branches.
If there's no node, we return 0, which is our base case.
4. Check if value exists in a tree
def valExists(root, val):
if not root:
return False
return root.val == val or valExists(root.left, val) or valExists(root.right, val)
The idea is to check if value exists in the root node, or the left branch or the right branch. If so, return True.
If we end up reaching a non-existent (NULL) node, we can safely return False (i.e. we can safely say the value doesn't exist in this Tree).
5. Invert a binary tree
Original:
1
/ \
3 2
/ \
5 4
Inverted:
1
/ \
2 3
/ \
4 5
def invert(root):
if not root:
return
root.left, root.right = root.right, root.left
invert(root.left)
invert(root.right)
We can invert a binary tree as follows:
- For the current node, swap values of its left and right branches.
- Repeat this for the node on the left and the node on the right.
- Stop when we hit an empty node.
6. Check if same tree
def sameTree(t1, t2):
if not t1 and not t2:
return True
if t1 and t2:
return t1.val == t2.val and sameTree(t1.left, t2.left) and sameTree(t1.right, t2.right)
return False
For both trees to be the same:
- both nodes need to exist.
- their root values need to match.
- their left branch and right branch needs to match (we can check this recursively).
If while checking we end up reaching a position where both branches are missing, we know we have hit the leaf nodes without any mismatch in values between two trees. Therefore we must return True when we hit leaf nodes.
In all other cases, we should return False.
7. Check if subtree of another tree
def isSubTree(root, subRoot):
if not root:
return False
return sameTree(root, subRoot) \
isSubTree(root.left, subRoot) \
isSubTree(root.right, subRoot)
To check if one tree is a subtree of another tree, simply re-use the sameTree function from above with the following logic:
- Is current node and subRoot the same tree? If so, return True.
- Otherwise recursively check its left and right branches.
- If we hit a leaf node, we should return False because we haven't found a True value for sameTree function so far.
8. Binary Search Tree (BST): Lowest Common Ancestor
In a BST, all the values on the left are strictly less than the current node while all the values on the right are strictly greater than the current node.
We can use this to our advantage in finding the lowest common ancestor as follows:
def LCAofBST(root, p, q):
if not root:
return False
if p.val < root.val and q.val < root.val:
return LCAofBST(root.left, p, q)
if p.val > root.val and q.val > root.val:
return LCAofBST(root.right, p, q)
return root
The logic is quite simple:
- If there's no root, return False.
- If both values are less than root, check the left side.
- If both values are greater than root, check the right side.
- In all other cases, we know p falls on the left and q falls on the right, this makes the current node the least common ancestor between the two, therefore we return root.
9. Lowest Common Ancestor in a Binary Tree
What if our tree is not a BST? Then finding the LCA is a little difficult but not impossible. Here's how to do it:
def lca(root, p, q):
if not root:
return None
if root.val == p or root.val == q:
return root
left_result = lca(root.left, p, q)
right_result = lca(root.right, p, q)
if left_result and right_result:
return root
if left_result:
return left_result
else:
return right_result
The logic here can be made simple with two key points:
- If the left branch and the right branch has p and q respectively, return the current node as LCA.
- If only one branch as p or q, return p or q respectively.
For example:
1
/ \
2 3
/ \
4 5
(A) Finding 4 and 5
- You hit 4, you return 4.
- You hit 5, you return 5.
- You hit 2, you get 4 and 5, so you return 2.
- You hit 3, you return None.
- You hit 1, you get 2 and None, you return 2.
(B) Finding 2 and 5
- You hit 2, you return 2.
- You hit 3, you return None.
- You hit 1, you get 2 and None, you return 2.
(C) Finding 2 and 3
- You hit 2, you return 2.
- You hit 3, you return 3.
- You hit 1, you get 2 and 3, you return 1.
10. Final tip
My final tip would be to visualize the tree and try to solve it for smaller problems as follows:
NULL
1
1
/ \
2 3
1
/
2
1
\
3
Imagine having a Null node, a single node, a node with two branches, a node with left branch, and finally a node with right branch.
How would you solve for these cases? This would help you make some progress on the given problem :)
Top comments (0)