LeetCode 222. Count Complete Tree Nodes满二叉树节点个数

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

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.

Example:

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

Output: 6

---------------------------------------------

O(n) complexity is easy. The following algorithm is O(logn)^2

# 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: TreeNode) -> int:
        def count_height(root):
            res = 0
            while (root):
                res += 1
                root = root.left
            return res
        
        def dfs(root):
            if (root == None):
                return 0
            
            hl,hr = count_height(root.left),count_height(root.right)
            if (hl == hr):
                return (1<<hl)+dfs(root.right)
            else:
                return (1<<hr)+dfs(root.left)

        return dfs(root)