Use Union Find to solve a hard problem

2076. Process Restricted Friend Requests

Are you sure you really master the UF data structure/algorithm? If you are, you can use this question to verify your skills of using UF. This is the 4th question of the LC weekly 267 contest. It is pretty fun.

Basic ideas

• basic UF
• + maintain members for each set (usually, we don not need do this for basic UF problems)
• + maintain a black list (hates list)
• + check whether two sets can stay with peace. If they can, then union them.

Code

The extra parts compared with standard UF are commented.

class UF{
    public:
    vector<int> roots;
    // extra variable compare with standard UF
    vector<unordered_set<int>> hates;
    vector<unordered_set<int>> members;
    UF(int n){
        roots.assign(n, 0);
        for (int i = 0; i < n; ++i){
            roots[i] = i;
            hates.push_back(unordered_set<int>{});
            members.push_back(unordered_set<int>{i});
        }
    }
    int find_set(int x){
        return roots[x] == x? x : roots[x] = find_set(roots[x]);
    }
    void union_set(int x, int y){
        int root_x = find_set(x), root_y = find_set(y);
        roots[root_y] = root_x;
        // extra steps compare with standard UF
        for (auto x: hates[root_y]){
            hates[root_x].insert(x);
        } 
        for (auto x: members[root_y]){
            members[root_x].insert(x);
        }
    }
    void update_hates(int x, int node){
        hates[x].insert(node);
    }
// extra member function compare with standard UF
    bool peace(int u, int v){
        auto root_u = find_set(u);
        auto root_v = find_set(v);
        return _peace(root_u, root_v) && _peace(root_v, root_u);
}
    private:
    bool _peace(int u, int node){
        for (auto x: members[node]){
            if (hates[u].find(x) != hates[u].end())return false;
        }
        return true;
    }
};
class Solution {
    public:
    vector<bool> friendRequests(int n, vector<vector<int>>& restrictions, vector<vector<int>>& requests) {
        UF uf(n);
        for (auto & r: restrictions){
            uf.update_hates(r[0], r[1]);
            uf.update_hates(r[1], r[0]);
        }
        vector<bool> ret;
        const int m = requests.size();
        for (int i = 0; i < m; ++i){
            auto req = requests[i];
            auto u = req[0], v = req[1];
            if (uf.peace(u, v)){
                ret.push_back(true);
                uf.union_set(u, v);
            }
            else ret.push_back(false);
        }
        return ret;
}
};

Take away

This is a good problem to verify your skills of UF. Hope you enjoy this problem.