Count Complete Tree Nodes (Interview Solution)

How to Solve Count Complete Tree Nodes

Count Complete Tree Nodes Introduction

The Count Complete Tree Nodes problem asks us to count the number of nodes in a complete binary tree. This problem requires careful consideration of what a complete binary tree is - where every level besides the last is filled in - and the traversal algorithms necessary to explore the tree efficiently.

Count Complete Tree Nodes Problem

Given the root of a complete binary tree, return the number of nodes in the tree.

For this problem a binary tree is considered “complete” if every level besides the last level is completely filled in.

Example Inputs and Outputs

Example 1

Input:

root = [1, 2]

Output: 2

Example 2

Input:

root = [1, 2, 3, 4, 5, 6]

Output: 6

Constraints

Count Complete Tree Nodes Solutions

Approach 1: Brute Force

When solving an algorithm, it is often best to start with the brute force solution, and from there you can optimize the solution to make it more efficient.

The brute force solution can be achieved by traversing the whole tree and counting the nodes. There are two main methods to traverse a tree, depth-first search (DFS) or breadth-first search (BFS), and each of those methods has pros/cons depending on the problem at hand. However, when the problem simply requires traversing the tree as a whole with no possibility of exiting the traversal early then either approach will work and you can proceed with whatever implementation you are more comfortable coding.

For this problem, let’s use DFS. DFS starts at the root of the tree and explores as deep as possible along each branch before backtracking. DFS can be implemented iteratively or recursively, and for most people both coding out and reading the recursive implementation is easier, so let’s go with the recursive implementation for the purposes of this problem.

# TreeNode definition
# class TreeNode:
#    def __init__(self, data):
#        self.data = data
#        self.left = None
#        self.right = None

def count_nodes_dfs(root: TreeNode):
   if root is None:
       return 0
   return count_nodes_dfs(root.left) + count_nodes_dfs(root.right) + 1

Time/Space Complexity

Approach 2: Binary Search

Every Level, Except Possibly The Last, Is Completely Filled

We know that in a complete binary tree, every level, except possibly the last, is completely filled. Let’s take the tree below where we have all levels completely filled. Can we think of a way to find the count of all nodes without traversing the whole tree?

Since all the levels are completely filled, we know that level 0 will have 1 node, level 1 will have 2 nodes, level 2 will have 4 nodes, and so on. So, the total number of nodes will be 2^0 + 2^1 + ... + 2^h where h is the height of the tree. This equals to 2^(h+1) - 1. Which means, in this case, we just need to find the height of the tree to calculate the number of nodes.

Last Level, All Nodes Are As Far Left As Possible

Let’s imagine we are playing a game. We have the complete binary tree below. We do not know which nodes in the last level are filled. Assume you have direct access to the last level’s nodes and you are allowed to reveal any of these nodes. The question is: can you identify which nodes in the last level are filled? Reveal as few nodes as you can. Remember, in a complete binary tree, all nodes are as far left as possible.

One solution is to reveal node A, then B, then C until you reach an empty node. This is a linear search. Can you think of a better way?

We can start by revealing the node in the middle, if it is empty, we therefore know that everything to its right is also empty. From there we continue the search on the left side.

Does this algorithm look familiar? It should – it is binary search! By using this method we are able to make our search faster, as the time complexity would be logarithmic (log(n)) instead of linear (n).

Binary Search and Complete Binary Tree

Now, if we are at the root of the tree, and we do not have direct access to the last level, how can we leverage the idea of binary search to find the rightmost node in the last level?

Let’s visualize that. We are at the root, node 1, our goal is to find the rightmost node in the last level. How can we decide if we should go left or right?

Well, we can check the height of the right node, and if it is equal to the height of the current node - 1, this means that there is a path in the right tree where it will take us to a node in the last level. In such a case, we can move to the right and ignore the left tree.

Specific to our example, the current node, node 1, has height = 3 and the right node, node 3, has height = 2. This means we can go right.

# TreeNode definition
# class TreeNode:
#    def __init__(self, data):
#        self.data = data
#        self.left = None
#        self.right = None

def calculate_height(node: TreeNode):
   if not node:
       return -1
   height = 0
   while node.left:
       node = node.left
       height += 1
   return height

def count_nodes(root: TreeNode):
   nodes_count = 0
   height = calculate_height(root)

while root:
       if calculate_height(root.right) == height - 1:
           nodes_count += pow(2, height)
           root = root.right
       else:
           nodes_count += pow(2, height - 1)
           root = root.left
       height -= 1
   return nodes_count

Time/Space Complexity