Count Complete Tree Nodes

Problem: Count Complete Tree Nodes

If we consider this tree as a heap and try to use an array to arrange and count nodes, the running time will exceed limit. The better solution is to use binary search. The general idea is that the last element is either in the right subtree or left subtree. We can easily judge that by observe the height of tree and subtree.

If the tree's height is h and right subtree's height is h-1, the last element is surely in the right subtree since the left subtree is full. Thus we can count root and number of left subtree and recursively do right part. Otherwise, the last element is in left tree and the right subtree's last level is missing. We can still count root and number of right subtree and recursively do left part.

Code in Python:

from collections import deque

class Solution(object):
    def countNodes(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        h = self.height(root)
        res = 0
        while root:
            if self.height(root.right) == h - 1:
                res += 2 ** h
                root = root.right
            else:
                res += 2 ** (h - 1)
                root = root.left
            h -= 1
        return res

    def height(self, root):
        return -1 if not root else 1 + self.height(root.left)