Concise Java solutions O(log(n)^2)


  • 280

    Main Solution - 572 ms

    class Solution {
        int height(TreeNode root) {
            return root == null ? -1 : 1 + height(root.left);
        }
        public int countNodes(TreeNode root) {
            int h = height(root);
            return h < 0 ? 0 :
                   height(root.right) == h-1 ? (1 << h) + countNodes(root.right)
                                             : (1 << h-1) + countNodes(root.left);
        }
    }
    

    Explanation

    The height of a tree can be found by just going left. Let a single node tree have height 0. Find the height h of the whole tree. If the whole tree is empty, i.e., has height -1, there are 0 nodes.

    Otherwise check whether the height of the right subtree is just one less than that of the whole tree, meaning left and right subtree have the same height.

    • If yes, then the last node on the last tree row is in the right subtree and the left subtree is a full tree of height h-1. So we take the 2^h-1 nodes of the left subtree plus the 1 root node plus recursively the number of nodes in the right subtree.
    • If no, then the last node on the last tree row is in the left subtree and the right subtree is a full tree of height h-2. So we take the 2^(h-1)-1 nodes of the right subtree plus the 1 root node plus recursively the number of nodes in the left subtree.

    Since I halve the tree in every recursive step, I have O(log(n)) steps. Finding a height costs O(log(n)). So overall O(log(n)^2).


    Iterative Version - 508 ms

    Here's an iterative version as well, with the benefit that I don't recompute h in every step.

    class Solution {
        int height(TreeNode root) {
            return root == null ? -1 : 1 + height(root.left);
        }
        public int countNodes(TreeNode root) {
            int nodes = 0, h = height(root);
            while (root != null) {
                if (height(root.right) == h - 1) {
                    nodes += 1 << h;
                    root = root.right;
                } else {
                    nodes += 1 << h-1;
                    root = root.left;
                }
                h--;
            }
            return nodes;
        }
    }
    

    A Different Solution - 544 ms

    Here's one based on victorlee's C++ solution.

    class Solution {
        public int countNodes(TreeNode root) {
            if (root == null)
                return 0;
            TreeNode left = root, right = root;
            int height = 0;
            while (right != null) {
                left = left.left;
                right = right.right;
                height++;
            }
            if (left == null)
                return (1 << height) - 1;
            return 1 + countNodes(root.left) + countNodes(root.right);
        }
    }
    

    Note that that's basically this:

    public int countNodes(TreeNode root) {
        if (root == null)
            return 0;
        return 1 + countNodes(root.left) + countNodes(root.right)
    

    That would be O(n). But... the actual solution has a gigantic optimization. It first walks all the way left and right to determine the height and whether it's a full tree, meaning the last row is full. If so, then the answer is just 2^height-1. And since always at least one of the two recursive calls is such a full tree, at least one of the two calls immediately stops. Again we have runtime O(log(n)^2).


  • 24
    D

    My code (check the left subtree is perfect or not), 76 ms even the time complexity is also o(height^2)

    class Solution {
    public:
        int countNodes(TreeNode* root) {
            if(!root) return 0;
            int num=1;
            TreeNode *curR(root->left), *curL(root->left);
            while(curR) // curR is the rightmost edge, which has a height equal to or less than the leftmost edge
            {
                curL = curL->left;
                curR = curR->right;
                num = num<<1;
            }
            return  num + ( (!curL)?countNodes(root->right):countNodes(root->left) );
        }
    };
    

  • 1
    B

    Hi @StefanPochmann thanks for your post.
    I understand the runtime as O(log(n)^2) for your main solution. However when applying the Master theorem to the recursion (https://en.wikipedia.org/wiki/Master_theorem),
    I got the time complexity O(log(n)), as the recurrence relation can be expressed as T(n) = 1 * T(n/ 2) + log(n).
    Do you know why the master theorem leads to a different runtime here?


  • 0

    The master theorem gives you O(log(n)^2) (actually theta, but I can't type that on my phone :-). You got the recurrence relation right, don't know what mistake you made. Which of the three cases did you use?


  • 0
    B

    I think according to the recurrence relation this is case 3.
    Because O(1) < f(n) = O(log(n)) < O(n) so c is a number between 0 and 1,
    a is 1 and b is 2, therefore c > logb(a) which maps to case 3.
    Is it correct? If so runtime of case 3 is O(f(n)) which is O(log(n))


  • 1

    No, log(n) is not Ω(nc) for c>0, as you can see here. You need to use case 2.


  • 0
    B

    right I see, I had some misunderstanding on master theory before. Thanks for your explanation Stefan! :)


  • 4
    H

    My Java code derived from your solution. Thanks for sharing!

    public class Solution {
    public int countNodes(TreeNode root) {
        if(root == null)
            return 0;
        TreeNode left = root.right, right = root.right;
        int height = 0;
        while(right!= null){
            left = left.left;
            right = right.right;
            height++;
        }
        if(left == null && right == null)
            return 1+(1 << height)-1 +countNodes(root.left);
        else
            return 1+(1 << height+1)-1 + countNodes(root.right);
    }
    

    }


  • 4

    The while-loop only ends with right==null, so you don't need to check that afterwards.


  • 0
    O

    Can't say more about reading your code, fantastic!


  • 0
    S

    Share my solution without repeating to calculate the subtree height again and again.
    kind of caching the sub problem result.

    class Solution {
    public:
        int countNodes(TreeNode* root) {
            return helper(root, -1, -1);
        }
    private:
        int helper(TreeNode* root, int leftH, int rightH){
            if(!root) return 0;
            TreeNode* cur;
            int lH, rH;
            if(-1==leftH){
                lH=0;
                cur=root;
                while(cur){
                    ++lH;
                    cur=cur->left;
                }
                leftH=lH;
            }
            
            if(-1==rightH){
                rH=0;
                cur=root;
                while(cur){
                    ++rH;
                    cur=cur->right;
                }
                rightH=rH;
            }
            
            if(rightH==leftH){
                return (1<<leftH)-1;
            }
            return 1+helper(root->left, leftH-1, -1)+helper(root->right, -1, rightH-1);
        }
    };

  • 4
    J

    Brilliant code!

    One minor issue. For empty tree, would it be even cleaner to use 0 instead of -1 as height? I tend to think any empty is a zero. I sound pretty dull. Do i? :-)

    Here is Java code :

    public class Solution {
    int height(TreeNode root) {
    	return root == null? 0 : 1 + height(root.left);
    }
    public int countNodes(TreeNode root) {
    	int h = height(root);
    	if (h == 0)	return 0;
        return height(root.right) == h-1? (1 << h-1) + countNodes(root.right)
        		: (1 << h - 2) + countNodes(root.left);
    } }

  • 0

    I think they're equally clean. I just define the height to be the length of the longest root-to-leaf path (counting steps), which for the single-node tree is zero.

    I don't remember, but I probably tried it your way as well and used mine because it allowed the rest of the code to be slightly shorter.


  • 0
    O
    public class Solution {
        public int countNodes(TreeNode root) {
            return helper(root, findDepth(root));
        }
        
        int helper(TreeNode root, int depth){
            if(root == null)
                return 0;
            int rightDepth = findDepth(root.right);
            return helper(rightDepth == depth-1?root.right:root.left, depth-1)+(1<<rightDepth);
        }
        
        int findDepth (TreeNode root){
            return root == null? 0:1+findDepth(root.left);
        }
    }
    

    A little improvement by passing current depth into next recursion.


  • -1
    R

    Excuse me a naive question, O(log(n)^2) is better than O(n)? Did I miss anything? The last solution is O(n), I don't see why it is slower than O(log(n)^2).


  • 0

    @Zhipeng_cpp Are you serious? Yes, O(log(n)^2) is far better than O(n).


  • 0
    R

    Oops, you are right, thx :D


  • 0
    B

    leetcode magician!


  • 0
    M

    Why is the right subtree one height shorter than the left subtree? ie why is it h-2 for counting the right subtree instead of h-1? Shouldn't it be symmetrical?


  • 0
    S

    I'm getting a Time Limit Exceeded with this code, which I believe is equivalent - what's wrong with this?

    public class Solution {
        public int countNodes(TreeNode root) {
            if (root == null) {
                return 0;
            }
            
            int h = height(root);
            if (h == 0) {
                return 0;
            }
            
            int rightHeight = height(root.right);
            if (rightHeight+1 == h) {
                return (int) Math.pow(2, h) + countNodes(root.right);
            } else {
                return (int) Math.pow(2, h-1) + countNodes(root.left);
            }
        }
        
        private int height(TreeNode root) {
            if (root == null) {
                return 0;
            }
            return 1 + height(root.left);
        }
    }
    

Log in to reply
 

Looks like your connection to LeetCode Discuss was lost, please wait while we try to reconnect.