Count Complete Tree Nodes

 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 }