LeetCode #222 Count Complete Tree Nodes

Posted on | Post modified | In , , , , , , , ,

Problem

Given a complete binary tree, count the number of nodes.

Example

1
2
3
4
5
6
7
8
Input: 
    1
   / \
  2   3
 / \  /
4  5 6

Output: 6

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.

Solution

Method: Binary Search

Time Complexity: O(log n)

Space Complexity:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def countNodes(self, root):
        """
        :type root: TreeNode
        :rtype: int
        """
        def find_depth(root):
            depth = 0
            while root:
                root = root.right
                depth += 1
            return depth

        count = 0
        start = root

        while start:
            left_depth = find_depth(start.left)
            right_depth = find_depth(start.right)
            #count += 1
            #count += 2 ** left_depth - 1
            count += 2 ** left_depth
            if left_depth > right_depth:
                start = start.right
            else:
                start = start.left

        return count