Description

Submission

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
    unordered_set<int> nonRoots;
    unordered_set<int> roots;
    unordered_set<int> used;
    unordered_map<int, TreeNode*> idx2root;
    TreeNode* root;

    void findNonRoot(TreeNode* root) {
        if(!root) return;
        nonRoots.insert(root->val);
        findNonRoot(root->left);
        findNonRoot(root->right);
    }

    bool buildTree(TreeNode*& root, int lower, int upper) {
        if(!root) return true;
        // cout << root->val << " " << lower << " " << upper << endl;
        if(root->val < lower || root->val > upper) return false;
        if(!root->left && !root->right) {
            // leaf node, find possible replacement
            if(!idx2root.count(root->val)) return true;
            if(used.count(root->val) && root != this->root) return false;
            used.insert(root->val);
            root = idx2root[root->val];
            return buildTree(root->left, lower, root->val - 1) 
                && buildTree(root->right, root->val + 1, upper);
            
        }
        return buildTree(root->left, lower, root->val - 1) 
                && buildTree(root->right, root->val + 1, upper);
    }
public:
    TreeNode* canMerge(vector<TreeNode*>& trees) {
        int n = trees.size();
        for(TreeNode* root: trees) {
            roots.insert(root->val);
            findNonRoot(root->left);
            findNonRoot(root->right);
            idx2root[root->val] = root;
        }

        vector<int> rootCandidates;

        for(auto x: roots) {
            if(nonRoots.count(x)) continue;
            rootCandidates.push_back(x);
        }

        if(rootCandidates.size() != 1) return nullptr;

        root = idx2root[rootCandidates[0]];
        used.insert(root->val);

        bool ok = buildTree(root, INT_MIN, INT_MAX);

        if(ok && used.size() == n) {
            return root;
        }
        return nullptr;
    }
};