Leetcode Problem - Lowest Common Ancestor

Problem link: https://leetcode.com/problems/lowest-common-ancestor-of-a-binary-tree/
Difficulty: Medium
Category: Trees

Step 1: Grasping the problem

This is a tree problem where a tree (by root) and two nodes (node1 and node2) within the tree is given. We will need to return LCA of node1 and node2.

Lowest Common Ancestor (LCA) of two nodes is defined as the node which is ancestor of both node1 and node2 and is at the lowest level of tree as possible. For example, in the below tree, node1 = 6 and node2 = 4. The LCA node cannot be 2 because node 2 is not an ancestor of node1 On the other hand, node 3 cannot be LCA node because there is another node (node 5) which is ancestor of both node1 and node2 and is in deeper level than 3. So the answer must be node 5.

problem-tree

Step 2: Add parent attribute

Currently only the root of the tree is given. That means we can only go from root to any child. But to solve this problem, we will need to be able to traverse from a child node to it’s parent. So the first task is to add a new data parent for each node so that we can easily traverse both down and up in the tree.

Unlike other statically typed programming languages, python has the powerful ability to create class attributes from outside of a class. We will use this to create the parent attribute right in each nodes.

To add an attribute to each node we need to traverse through all the nodes once. Any tree traversal algorithm can be used here. We are using simple DFS here. Go ahead and change the stack to a queue to do a BFS traversal.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
root.parent = None
# Start a DFS and assign parent to each node
stack = {root}
while stack:
    node = stack.pop()
    if node.left:
        node.left.parent = node
        stack.append(node.left)
    if node.right:
        node.right.parent = node
        stack.append(node.right)

Step 3: Store the path of `node1`

Now we need to store the path of node node1 and then look where the path conflicts with node node2’s path. node1 and node2 are interchangeable in this case so you can store the path of node2 and look up which node in node1 path conflicts with it.

As we already have the node node1 in hand, with the help of parent attribute we just created, we will traverse up in the tree till root and store all the nodes in the path.

path-store

This stored list of nodes will only be used for lookup of a node. So instead of using list as the data structure, we can use a hashset implementation set which is similar to unordered_set<int> in C++ and efficient in key lookups.

1
2
3
4
5
# Trace the path of node node1
node1_path = set()
while node1:
    node1_path.append(node1)
    node1 = node1.parent

Step 4: Search for first existent node

Now it is needed to do the almost same thing with node node2. Except for storing the nodes of the path, we will loop up the node in node1_path. The first node we find will be our answer!

lookup

1
2
3
4
5
# Search for first node of node2's parent in node1_path
while node2:
    if node2 in node1_path:
        return node2
    node2 = node2.parent

This is it. A good practice though is to return a value at the end of the function (just in case one cannot find and claim the code to be broken).

1
return None # A good practice

So the overall time complexity of the function is O(n) because there is nothing but traversal we are doing. The space complexity is also O(n) because we are storing the path of one node (node1 or node2). There are algorithms where you can make the space complexity to O(1) in the cost of more time. Below is the complete code of the solution.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# Definition for a binary tree node.
class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

class Solution:
    def lowestCommonAncestor(self, root: TreeNode, node1: TreeNode, node2: TreeNode) -> TreeNode:
        root.parent = None

        # Start a DFS and assign parent to each node
        stack = [root]
        while stack:
            node = stack.pop()
            if node.left:
                node.left.parent = node
                stack.append(node.left)
            if node.right:
                node.right.parent = node
                stack.append(node.right)

        # Trace the path of node node1
        node1_path = set()
        while node1:
            node1_path.add(node1)
            node1 = node1.parent
        
        # Search for first node of node2's parent in p_path
        while node2:
            if node2 in node1_path:
                return node2
            node2 = node2.parent
        
        return None # A good practice

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    TreeNode* lowestCommonAncestor(TreeNode* root, TreeNode* node1, TreeNode* node2) {
        // Start a DFS and assign to each node
        stack<TreeNode*> st;
        st.push(root);
        map<TreeNode*, TreeNode*> parents;
        while(!st.empty()) {
            auto node = st.top();
            st.pop();
            if(node->left != nullptr) {
                parents[node->left] = node;
                st.push(node->left);
            }
            if(node->right != nullptr) {
                parents[node->right] = node;
                st.push(node->right);
            }
        }

        // Trace the path of node node1
        set<TreeNode*> node1Path;
        while(node1 != nullptr) {
            node1Path.insert(node1);
            node1 = parents.find(node1) != parents.end() ? parents[node1] : nullptr;
        }

        // Search for first node of node2's parent in node1Path
        while(node2 != nullptr) {
            if(node1Path.find(node2) != node1Path.end()) {
                return node2;
            }
            node2 = parents.find(node2) != parents.end() ? parents[node2] : nullptr;
        }
        return nullptr; // Must return in C++
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */

class Solution {
    public TreeNode lowestCommonAncestor(TreeNode root, TreeNode node1, TreeNode node2) {

        // Start a DFS and assign parent to each node
        Stack<TreeNode> stack = new Stack<>();
        stack.push(root);
        Map<TreeNode, TreeNode> parents = new HashMap<>();
        while(!stack.isEmpty()) {
            TreeNode node = stack.pop();
            if(node.left != null) {
                parents.put(node.left, node);
                stack.push(node.left);
            }
            if(node.right != null) {
                parents.put(node.right, node);
                stack.push(node.right);
            }
        }

        // Trace the path of node node1
        Set<TreeNode> node1Path = new HashSet<TreeNode>();
        while(node1 != null) {
            node1Path.add(node1);
            node1 = parents.getOrDefault(node1, null);
        }

        // Search for first node of node2's parent in node1_path
        while(node2 != null) {
            if(node1Path.contains(node2))
                return node2;
            node2 = parents.getOrDefault(node2, null);
        }
        return null; // Must return in Java
    }
}