[LeetCode] Problem 222 - Count Complete Tree Nodes

Posted on | In ,

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 2^h nodes inclusive at the last level h.

Example

Input: 

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

Output: 6

Code

1
2
3
4
5
6
public class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;
    TreeNode(int x) { val = x; }
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
public int countNodes(TreeNode root) {
    int leftHeight = getLeftHeight(root);
    int rightHeight = getRightHeight(root);

    if (leftHeight == rightHeight)
        return (int) Math.pow(2, leftHeight) - 1;
    else
        return countNodes(root.left) + countNodes(root.right) + 1;
}

private int getLeftHeight(TreeNode root) {
    int height = 0;

    while (root != null) {
        height++;
        root = root.left;
    }

    return height;
}

private int getRightHeight(TreeNode root) {
    int height = 0;

    while (root != null) {
        height++;
        root = root.right;
    }

    return height;
}