1 /**
2 * Definition for a binary tree node.
3 * public class TreeNode {
4 * int val;
5 * TreeNode left;
6 * TreeNode right;
7 * TreeNode(int x) { val = x; }
8 * }
9 */
10 public class Solution {
11
12 int height(TreeNode root){
13 if(root == null) return 0;
14
15 return 1 + height(root.left);
16 }
17
18
19 int leaves = 0;
20 boolean stop = false;
21
22 void countLeaves(TreeNode root, int heightToLeaf){
23
24 if(root == null) return;
25
26 if(stop) return;
27
28 if(heightToLeaf == 2) {
29
30 if(root.right != null){
31 leaves += 2;
32 } else {
33
34 // at lease one is null
35 stop = true;
36
37 if(root.left != null) {
38 leaves += 1;
39 }
40 }
41
42
43 return;
44 }
45
46 countLeaves(root.left, heightToLeaf - 1);
47 countLeaves(root.right, heightToLeaf - 1);
48 }
49
50 int perfectTreeNodeCount(int height){
51 if(height == 0) return 0;
52 if(height == 1) return 1;
53
54 return (int)Math.pow(2, height) - 1;
55 }
56
57 public int countNodes(TreeNode root) {
58
59 int h = height(root);
60
61 countLeaves(root, h);
62
63 if(!stop){
64 // perfect tree
65 return perfectTreeNodeCount(h);
66 }
67
68 return perfectTreeNodeCount(h - 1) + leaves;
69 }
70 }