Count Complete Tree Nodes | Solution (Explained)

LeetCode 222. Count Complete Tree Nodes is a medium problem if you go by the naive solution. To upgrade the time complexity from O(N) though, we need some more tricks up our sleeves. In this blog I explore both the solutions. Problem Link.

0. Count Complete Tree Nodes: Naïve Solution

0.0 Code

class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        def dfs(node):
            if not node: return 0
            return dfs(node.left) + dfs(node.right) + 1
        
        return dfs(root)

0.1. Complexity Analysis

• Time: O(N) to visit all the nodes
• Space: O(LogN), the height of the recursion call stack, the height of binary tree.

1. Observations

A bit sick right now, will add the explanations later. Put a comment if I forget to put by the end of week 😛

2. Count Complete Tree Nodes: Optimized Solution

2.0. Code

class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        def dfs(node):
            if not node: return 0
            
            lh = get_height(node, 0)
            rh = get_height(node, 1)
            if lh == rh: return ((1 << lh) - 1)
            
            return dfs(node.left) + dfs(node.right) + 1
        
        def get_height(node, i):
            count = 1
            
            if i == 0:
                while (node := node.left): count += 1
            else:
                while (node := node.right): count += 1
            
            return count
            
        return dfs(root)

2.1. Complexity Analysis

• Time: O(LogN * LogN). Each time, we either choose to go for the lh == rh option OR the standard route. Either way we compute the depth in O(LogN) time. This decision takes place for each height and not each node. Since we would bypass the computation a lot of times – due to the lh == rh being followed. So,
  • depth computation: O(LogN)
  • #times depth is computed: O(LogN)
• Space: O(LogN) for the recursion call stack height.

Tanishq Chaudhary

Producing high-quality intuitive explanations of interview problems. Currently covering LeetCode and InterviewBit.