Leetcode 222. Count Complete Tree Nodes

https://leetcode.com/problems/count-complete-tree-nodes/

If use recursion ignoring the feature of completed binary tree, the run time will get TLE.

First we should get the level of this tree. Because the left node will be there first, we can go along the left and gain the level. Then we must decide how many nodes are in this level. This is not very easy, and we can use binary search to find the most right node.

Solution:

# T:O((logn)^2) S:O(1)
class Solution(object):
    def countNodes(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        if root is None: return 0
        
        node, level = root, 0
        while node.left is not None:
            node = node.left
            level += 1
        
        # Binary search.
        left, right = 2 ** level, 2 ** (level + 1)
        while left < right:
            mid = left + (right - left) / 2
            if not self.exist(root, mid):
                right = mid
            else:
                left = mid + 1
                
        return left - 1
    
    # Check if the nth node exist.
    def exist(self, root, n):
        k = 1
        while k <= n:
            k <<= 1
        k >>= 2
        
        node = root
        while k > 0:
            if (n & k) == 0:
                node = node.left
            else:
                node = node.right
            k >>= 1
        return node is not None

Run Time: 188 ms