Use Union Find to solve a hard problem

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

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

--

--

--

Data Scientist/MLE/SWE @takemobi

Love podcasts or audiobooks? Learn on the go with our new app.

Recommended from Medium

Many of the most popular YouTube video reviews are the same:

Release multi-target Rust applications with GitLab CI

Creating Partitioned tables in Big Query

Generic Interfaces - .Net

Laravel Tutorial #4: Database and Models

My “Almost Ultimate” Work From Home Setup as a Cloud Architect

SCD Implementation with Temporal Tables in Power BI

SCD Implementation with Temporal Tables in Power BI

Three Ways of Achieving Rapid Application Development in IoT

Get the Medium app

A button that says 'Download on the App Store', and if clicked it will lead you to the iOS App store
A button that says 'Get it on, Google Play', and if clicked it will lead you to the Google Play store
Jimmy Shen

Jimmy Shen

Data Scientist/MLE/SWE @takemobi

More from Medium

Introduction to Binary Search Algorithm

Contiguous Subarrays

1 — Two Sum