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

countCompleteTreeNodes.py

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

# O(h^2)
# Tree height = h
# Range of possible answers, r = 2^h
# Complexity = h * log2(r)
#            = h * log2(2^h)
#            = h * h
#            = h^2
class Solution:
    def countNodes(self, root: TreeNode) -> int:
        
        if root is None: return 0

        h, curr = 0, root.left
        while(curr):
            h+=1
            curr = curr.left
        
        result, l, r = 1, 1, pow(2, h)
        
        while(l <= r):
            mid = (l+r)//2
            
            if self.nodeExistsAt(root, mid, h):
                result = max(result, mid)
                l = mid+1
            else:
                result = min(result, mid)
                r = mid-1
        
        return result + (1 << h)-1
            
    def nodeExistsAt(self, root, target, maxHeight):
        
        l, r, depth = 1, pow(2, maxHeight), 0

        while(root is not None):
            mid = (l+r)//2
            if target > mid:
                root = root.right
                l = mid+1
            else:
                root = root.left
                r = mid-1
            if root is not None: depth+=1
        return depth == maxHeight