[LeetCode] 222. Count Complete Tree Nodes_Medium tag: Binary search

 

 

Ideas:

1. T: O(n), S: O(lgn), depth of the tree

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def countNodes(self, root: Optional[TreeNode]) -> int:
        return self.countNodes(root.left) + self.countNodes(root.right) + 1 if root else 0

 

 

2. T: O: (lg n) ^ 2

1. 找到tree 的depth      O(lg n)

2. 建一个helper function来看index 是否exist， O (lg n)

3. 找last index that node exists.

4. 然后将最后一level的nodes + 之前所有node的数目，最后返回结果即可

 

Code:

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    # get the depth of the tree
    def countDepth(self, node: TreeNode) -> int:
        d = 0
        while node:
            d += 1
            node = node.left
        return d

    # 根据idx的位置来看是左边还是右边
    def exists(self, idx: int, d : int, node: TreeNode) -> bool:
        l, r = 0, 2 ** (d-1) - 1
        for _ in range(d - 1):
            mid = l + (r - l)//2
            if idx > mid:
                node = node.right
                l = mid
            else:
                node = node.left
                r = mid
        return node is not None
    

    # 画图来看，把最后一行弄成index in [0, 2 **(d - 1) - 1], 用二分法来找last index exists.
    def countNodes(self, root: Optional[TreeNode]) -> int:
        if not root: return 0
        d = self.countDepth(root)

        l, r = 0, 2 ** (d-1) - 1
        # 每一次二分，为lg(n) , 再确认exist不exist再一次lg(n), so total T: O(lg(n) ^ 2) 
        # n is the total nodes
        while l  + 1 < r:
            mid = l + (r - l)//2
            if self.exists(mid, d, root):
                l = mid
            else:
                r = mid
        
        # last level nodes = last existing index + 1 
        last_level = r +1 if self.exists(r, d, root) else l +1
        return (2 ** (d - 1) - 1) + last_level