My Accepted Java Iterative Solution


  • 1
    W
    public int countUnivalSubtrees(TreeNode root) {
            if(root==null) return 0;
            if(root.left==null && root.right==null) return 1;
            Stack<TreeNode> stack=new Stack<TreeNode>();
            Stack<TreeNode> stackR=new Stack<TreeNode>();
            Set<TreeNode> set=new HashSet<TreeNode>();
            stack.push(root);
            while(!stack.isEmpty()){
                TreeNode node=stack.pop();
                if(node.left==null && node.right==null) set.add(node);
                else if(!stackR.isEmpty() && stackR.peek()==node){
                    stackR.pop();
                    if(node.left==null && set.contains(node.right) && node.right.val==node.val) set.add(node);
                    else if(node.right==null && set.contains(node.left) && node.left.val==node.val) set.add(node);
                    else if(node.left!=null && node.right!=null && node.left.val==node.right.val
                        && node.left.val==node.val) set.add(node);
                }else{
                    stack.push(node);
                    stackR.push(node);
                    if(node.right!=null) stack.push(node.right);
                    if(node.left!=null) stack.push(node.left);
                }
            }
            return set.size();
        }

  • 0
    G

    Nice solution, however, I think there need to change a little bit in your code:

    change:

    else if(node.left!=null && node.right!=null && node.left.val==node.right.val
    && node.left.val==node.val) set.add(node);

    to:

    else if(node.left!=null && node.right!=null && set.contains(node.left) && set.contains(node.right) && node.left.val==node.right.val
    && node.left.val==node.val) set.add(node);

    Then, it can deal with the following case:

           1
          /
         2
        / \
       2   2
      / \
     2   3

Log in to reply
 

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