[Leetcode]222. Count Complete Tree Nodes

Given a complete binary tree, count the number of nodes.
Note: 
Definition of a complete binary tree from Wikipedia:
In a complete binary tree every level, except possibly the last, is completely filled, and all nodes in the last level are as far left as possible. It can have between 1 and 2h nodes inclusive at the last level h.
Example:
Input: 
    1 
   / \ 
  2   3 
 / \ / 
4 5 6 
Output: 6

Brute force approach – O(n)

Well, there’s a binary tree, one of the most straightforward solution is to recurse it! You can traverse the tree and count the node one by one. Check Tree Traversal – Recursively & Iteratively – Preorder, Inorder, Postorder to find different ways for tree traversal! Here’s a BFS solution and the time complexity is O(n):

    def countNodes(self, root: TreeNode) -> int:
        ret = 0
        q = [root]
        while q:
            node = q.pop(0)
            if node:
                ret += 1
                q.append(node.left)
                q.append(node.right)
        return ret

Optimised approach 1 – O((logn)^2)

It explicitly says that it’s a complete binary tree, so taking advantage of it would offer an optimized solution! These ideas are inspired from @StefanPochmann.

If we calculate the height h_left along all the left children and h_right along all the right children, there are only two cases:

  • h_left == h_right, so the tree is a full binary tree, and the nodes are 2^h-1.
  • h_left != h_right, so the tree is not a full binary tree, and the nodes are 1 + count(node.left)+count(node.right)

The way that we use to find the heights of left and right subtree is to iteratively look for left‘s left child and right‘s right child until right is None. Then if left is None, it is full tree. The code is as follows:

class Solution:
    def helper(self, node):
        if not node:
            return -1
        else:
            return 1+self.helper(node.left)
    def countNodes(self, root: TreeNode) -> int:
        if not root:
            return 0
        left = root
        right = root
        h = 0
        while right:
            left = left.left
            right = right.right
            h += 1
        if left is None:
            return (1<<h) - 1
        else:
            return 1 + self.countNodes(root.left) + self.countNodes(root.right)

The counting height is O(logn) and it iterates vertically through the tree, therefore the time complexity is O((logn)^2

optimised approach 2 – O((logn)^2))

Similarly, the height of a complete binary tree could be found by traversing all the way left. When we compare the heights of the left subtree and the right subtree, there are two cases:

  • h_left == h_right, so the last node is on the right subtree, and the nodes on the left subtree (a full tree) are 2^{h\_left}-1.
  • h_left != h_right, so the last node is on the left subtree, and the nodes on the right subtree are 2^{h\_left-1}-1, as the right subtree is a h_left-1 full tree.
class Solution:
    def helper(self, node):
        if not node:
            return -1
        else:
            return 1+self.helper(node.left)
    def countNodes(self, root: TreeNode) -> int:
        h = self.helper(root)
        if h < 0:
            return 0
        else:
            if self.helper(root.right) == h-1:
                return (1 << h) + self.countNodes(root.right)
            else:
                return self.countNodes(root.left) + (1 << h-1)

Reference

https://leetcode.com/problems/count-complete-tree-nodes/discuss/61958/Concise-Java-solutions-O(log(n)2)

Subscribe
Notify of
guest
0 Comments
Inline Feedbacks
View all comments