C++ 4 lines O(n) unordered_multiset


  • 3
    V

    I am traversing the tree and inserting sums of subtrees in the hash set. The tree can be partitioned if there is a subtree sum that is half of the whole tree sum. Note that I am using multiset as it shortens the code a bit and also to check for one-node tree.

    int tSum(TreeNode* r, unordered_multiset<long long>& s) {
        return r == nullptr ? 0 : *s.insert(r->val + tSum(r->left, s) + tSum(r->right, s));
    }
    bool checkEqualTree(TreeNode* root) {
        unordered_multiset<long long> sums;
        auto sum = tSum(root, sums);
        return sum % 2 == 0 && sums.size() > 1 && sums.find(sum / 2) != sums.end();
    }
    

  • 3

    @votrubac Mine is similar but simpler:

    bool checkEqualTree(TreeNode* r) {
        unordered_set<int> s;
        int p = sum(r, s, true);
        return !(p % 2) && s.count(p / 2) > 0;
    }
    
    int sum(TreeNode *r, unordered_set<int> &s, bool isroot) {
        if (!r) return 0;
        auto v = r->val + sum(r->left, s, false) + sum(r->right, s, false);
        return isroot ? v : *s.insert(v).first;
    }
    

    Fixed failure for case [0,1,-1]


  • 0
    V

    @mzchen good idea to use a single insert instead of two. It makes it cleaner. I updated my solution. BTW, your code won't work for an empty tree (just saying, it does not matter for this problem).


  • 0
    T
    /**
     * Definition for a binary tree node.
     * public class TreeNode {
     *     int val;
     *     TreeNode left;
     *     TreeNode right;
     *     TreeNode(int x) { val = x; }
     * }
     */
    class Solution {
        int max=Integer.MIN_VALUE;
        public boolean checkEqualTree(TreeNode root) {
            Map<Integer,Integer> map = new HashMap<>();
            if(root==null || (root.left==null&&root.right==null)){
                return false;
            }
            
            int sum = traverse(root,map);
            if(sum%2!=0){
                return false;
            }
            if(sum==0){
                return map.get(0)>1;
            }
            return map.containsKey(sum/2);
        }
        
        public int traverse(TreeNode root,Map<Integer,Integer> map){
            if(root==null){
                return 0;
            }
            
            int sum=root.val;
            sum+=traverse(root.left,map)+traverse(root.right,map);
            map.put(sum,map.getOrDefault(sum,0)+1);
            max=Math.max(max,sum);
            return sum;
        }
    }
    
    

  • 0

    I find your counting misleading... you added six lines.


  • 1

    It fails for example [0,1,-1] (you return true).


  • 0
    T

    @StefanPochmann you are right! we should use map instead of set here to traverse


  • 0

    @StefanPochmann Wrong output was given for the case [0,1,-1] by the judge system. c++ lang.


  • 0
    C

    Like others said, your implementation has a bug so it fails on cases where all nodes added to zero. My fix is to separate root to root->left and root->right. This will prevent the total sum to enter the set. By the way, you don't need a multiset, a unordered_set will do fine. See here for a bug-free version : https://discuss.leetcode.com/topic/100425/c-with-unordered_set


  • 0

    @votrubac To take care of the sum == 0 case, just need to check sums.count(sum/2) > (sum==0) to make sure 0 counted at least twice.

        bool checkEqualTree(TreeNode* r) {
            int sum = getSum(r);
            return sum%2==0 && sums.count(sum/2) > (sum==0);
        }
        
        unordered_multiset<int> sums;    
        int getSum(TreeNode* x) {
            return x? *sums.insert(x->val + getSum(x->left) + getSum(x->right)) : 0;
        }
    

Log in to reply
 

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