222. Count Complete Tree Nodes

222. Count Complete Tree Nodes.java

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public int countNodes(TreeNode root) {
	if (root==null) return 0;
	if (root.left==null) return 1;
	int height = 0;
    int nodesSum = 0;
	TreeNode curr = root;
    while(curr.left!=null) {
    	nodesSum += (1<<height);
    	height++;
    	curr = curr.left;
    }
    System.out.println(height);
    return nodesSum + countLastLevel(root, height);
}

    private int countLastLevel(TreeNode root, int height) {
    	if(height==1) 
    		if (root.right!=null) return 2;
    		else if (root.left!=null) return 1;
    		else return 0;
    	TreeNode midNode = root.left;
    	int currHeight = 1;
    	while(currHeight<height) {
    		currHeight++;
    		midNode = midNode.right;
    	}
    	if (midNode==null) return countLastLevel(root.left, height-1);
    	else return (1<<(height-1)) + countLastLevel(root.right, height-1);
    }
}