# 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:
        if not root:
            return 0

        def _left_height(node: Optional[TreeNode]) -> int:
            ans = 0
            while node:
                ans += 1
                node = node.left
            return ans

        def _right_height(node: Optional[TreeNode]) -> int:
            ans = 0
            while node:
                ans += 1
                node = node.right
            return ans

        def count_nodes(
            node: Optional[TreeNode],
            left_height: Optional[int],
            right_height: Optional[int],
        ) -> int:
            if node is None:
                return 0

            if left_height is None:
                left_height = 1 + _left_height(node.left)
            if right_height is None:
                right_height = 1 + _right_height(node.right)

            if left_height == right_height:
                return (2 ** left_height) - 1

            return (
                1
                + count_nodes(node.left, left_height - 1, None)
                + count_nodes(node.right, None, right_height - 1)
            )

        return count_nodes(root, None, None)