Tree Problems Made Easy

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:

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:

(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.

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:

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 :)

Blog