Taro Logo

Count Complete Tree Nodes

Easy
Google logo
Google
1 view
Topics:
TreesRecursion

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

According to Wikipedia, every level, except possibly the last, is completely filled in a complete binary tree, and all nodes in the last level are as far left as possible. It can have between 1 and 2^h nodes inclusive at the last level h.

Design an algorithm that runs in less than O(n) time complexity.

For example:

  • If the input tree is [1,2,3,4,5,6], the output should be 6.
  • If the input tree is empty [], the output should be 0.
  • If the input tree is just a single node [1], the output should be 1.

Write a function to efficiently count the nodes in a complete binary tree.

Solution


Naive Approach

The most straightforward way to count the number of nodes in a binary tree is to perform a simple traversal (e.g., depth-first search or breadth-first search) and increment a counter for each node visited. This approach visits every node in the tree.

def count_nodes_naive(root):
    if not root:
        return 0
    return 1 + count_nodes_naive(root.left) + count_nodes_naive(root.right)

Time Complexity: O(n), where n is the number of nodes in the tree. Space Complexity: O(h) in the average case, where h is the height of the tree, due to the recursion stack. In the worst case (skewed tree), it can be O(n).

Optimal Solution

Leveraging the fact that the tree is a complete binary tree, we can significantly improve upon the naive approach. The main idea is to determine the height of the left and right subtrees. If they are equal, the left subtree is a full binary tree, and we can calculate its size directly using the formula 2^h - 1, where h is the height. Otherwise, we recursively count nodes in both subtrees.

def count_nodes_optimal(root):
    if not root:
        return 0

    left_height = 0
    right_height = 0

    left_node = root.left
    right_node = root.right

    while left_node:
        left_height += 1
        left_node = left_node.left

    while right_node:
        right_height += 1
        right_node = right_node.right

    if left_height == right_height:
        return (1 << (left_height + 1)) - 1  # 2^(h+1) - 1
    else:
        return 1 + count_nodes_optimal(root.left) + count_nodes_optimal(root.right)

Explanation:

  1. Base Case: If the root is None, the tree is empty, so return 0.
  2. Calculate Heights: Determine the heights of the left and right subtrees by traversing down the leftmost and rightmost branches, respectively.
  3. Check for Full Tree: If the heights are equal, the left subtree is a full binary tree. The number of nodes in a full binary tree of height h is 2^(h+1) - 1. We add 1 to h because height is edge count, and we need level count.
  4. Recursive Call: If the heights are not equal, recursively count the nodes in both subtrees and add 1 (for the root node).

Edge Cases:

  • Empty Tree: If the root is None, the function immediately returns 0.
  • Single Node: If the tree consists of only the root node, the left and right heights will be 0, so left_height == right_height will be true, and the function will return (1 << 1) - 1 = 1.

Time Complexity: O((log n)^2), where n is the number of nodes in the tree. Calculating the height of the tree takes O(log n) time, and in the worst case, we might have to traverse down to leaf nodes O(log n) times. Space Complexity: O(log n) due to the recursion stack. In each recursive call, the depth of the stack increases by one, and in the worst case, the maximum depth is proportional to the height of the tree, which is O(log n).