Leetcode 222: Count Complete Tree Nodes

 

Code

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

class Solution(object):
    def countNodes(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        if root is None:
            return 0        

        height_left, height_right = -1, -1
        
        node = root
        while node:
            height_left += 1 
            node = node.left
       
        node = root
        while node:
            height_right += 1
            node = node.right
       
        if height_left == height_right:
            return 2**(height_left+1) - 1
        else:
            return 1 + self.countNodes(root.left) + self.countNodes(root.right)

 

Idea

height_left is the number of edges between root and its leftmost leaf. height_right is the number of edges between root and its rightmost leaf. If they are the same mean, that means the last level is full. Therefore we can return 2**(height_left+1)-1 by the property of binary tree. If height_left does not equal to height_right, then we have to recursively count the nodes in the left and right subtree.

The time complexity if O((logn)^2). This is because you need O(logn) to calculate height_left/height_right. After that, the worst case is that you need to recursively call self.countNodes(root.left) and self.countNodes(root.right). This happens no more than O(logn) times.

 

Reference:

https://discuss.leetcode.com/topic/15515/easy-short-c-recursive-solution