Count Complete Tree Nodes

Problem Id: 222 Difficulty: Medium Tag: Binary Search Tag: Tree


Intuition

The brute force solution of this problem is to do a traversal through the tree and could the total number of nodes. But in this artical I will give an much faster way.

For a complete tree, all nodes except for the last row is filled. So the problem becomes finding how many nodes are there in the last row. And since the left most nodes on the last row is filled, we need to find the last node on the last row.

So we could use binary search. To judge if the ith node is empty on the last level, we need to go down of the tree through one path and the time complexity is O(lg(n)). We have totaly n // 2 nodes on the last row, so to do a binary search, it would totally take O(lg(n) * lg(n)) time. And it is faster than O(n).

Solution


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    @staticmethod
    def _depth(root):
        depth = 0
        while root is not None:
            depth += 1
            root = root.left
        return depth

    def _ith_depth(self, node, i):
        mod = 2 ** (self.depth - 2)
        depth = 1
        while node:
            depth += 1
            if i < mod:
                node = node.left
            elif i >= mod:
                i -= mod
                node = node.right
            mod //= 2
        return depth

    def countNodes(self, root: TreeNode) -> int:
        self.depth = self._depth(root)
        lo = 0
        hi = 2 ** (self.depth - 1)
        while hi - lo > 1:
            mid = (hi + lo) // 2
            if self._ith_depth(root, mid) == self.depth:
                hi = mid
            else:
                lo = mid
        left = (2 ** (self.depth - 1)) - hi
        return int((2 ** (self.depth)) - left - 1)