Count the number of nodes of complete binary tree

Count the number of nodes of complete binary tree

• Time: O(N) Space: O(1)

public int countNodes(TreeNode root) {
    int left = getLeftHeight(root);
    int right = getRightHeight(root);
    if (left == right) return (1 << left) - 1;
    else return countNodes(root.left) 
              + countNodes(root.right) + 1;
}
private int getLeftHeight(TreeNode root) {
    int i = 0;
    while (root != null) {
        root = root.left;
        i++;
    }
    return i;
}
private int getRightHeight(TreeNode root) {
    int i = 0;
    while (root != null) {
        root = root.right;
        i++;
    }
    return i;
}