The hardest union find problem ever
I have been playing in LC for a while and have solved quite a lot of problems by using union-find. If you want to see some of my previous posts, please check my personal website “software engineer” part’s union-find topic.
Yesterday, I encountered the hardest union find the problem so far and it is quite interesting due to the:
- Early prune the union-find graph from the beginning.
- Build the union-find graph gradually to collect the results.
Problem
Explanation
A good explanation can be found in [2.]
Key ideas can be summarized:
For a specified value
- We can use union find to organize all the nodes with value ≤ this specified value.
- For the connected component, the number of the valid paths is equal to n choose 2 where n is the number of nodes in this connected graph, whose value == the specified value. The reason is in a tree graph, there is only one unique path to connect any two nodes. Details can be found here
In order to speed up the union-find process, we can sort the values and buildup the graph by using union-find gradually, which can make sure:
- The graph defined by smaller values can all be used to make sure the correctness of the current explored value.
- The larger value nodes will be added when only when needed.
Naive backtrack solution (TLE)
We can solve the problem by using a naive backtrack solution:
From each start node, backtrack all valid paths to harvest the valid one.
This solution can pass most of the cases and due to high time complexity got TLE as expected.
class Solution {
public:
unordered_map<int, vector<int>> graph;
void build_graph(vector<vector<int>>& edges, vector<int>&vals ){
for (auto &e: edges){
auto u = e[0], v = e[1];
graph[v].push_back(u);
graph[u].push_back(v);
}
}
void backtrack(int i, int begin_idx, int val, vector<int>&vals, vector<int> &seen, int& paths, vector<int>& curr_path){
if (vals[i] == val && curr_path.size() > 1 && i > begin_idx){
paths++;
}
for (auto j: graph[i]){
if (seen[j] || vals[j] > val)continue;
curr_path.push_back(j);
seen[j] = 1;
backtrack(j,begin_idx, val, vals, seen, paths, curr_path);
seen[j] = 0;
curr_path.pop_back();
}
}
int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
const int n = vals.size();
int paths = 0;
build_graph(edges, vals);
for (int i = 0; i < n; ++i){
vector<int> seen(n, 0);
vector<int> curr_path;
seen[i] = 1;
curr_path.push_back(i);
int val = vals[i];
backtrack(i, i, val, vals, seen, paths, curr_path);
}
return paths + n;
}
};Naive union find (TLE)
class UF{
public:
vector<int> roots;
vector<int> num_val;
int val;
UF(int n, int val, vector<int> vals): val(val){
roots.assign(n, 0);
num_val.assign(n, 0);
for (int i = 0; i < n; ++i){
roots[i] = i;
num_val[i] = vals[i] == val;
}
}
int find_set(int x){
return roots[x] == x?x: roots[x] = find_set(roots[x]);
}
void union_set(int x, int y){
auto root_x = find_set(x);
auto root_y = find_set(y);
if (root_x != root_y){
roots[root_y] = root_x;
num_val[root_x] += num_val[root_y];
}
}
};
int nchoose2(int n ){
return n*(n-1)/2;
}
class Solution {
public:
int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
unordered_set<int> unique_vals(vals.begin(), vals.end());
int ret = 0;
const int n = vals.size();
// based on each unique val, build graph by using union find
for (auto val: unique_vals){
UF uf(n, val, vals);
for (auto e:edges){
int u = e[0], v = e[1];
if (vals[u] <= val && vals[v] <= val){
uf.union_set(u, v);
}
}
unordered_set<int> explored_roots;
for (int i = 0; i < n; ++i){
auto root = uf.find_set(i);
if (explored_roots.find(root) == explored_roots.end()){
explored_roots.insert(root);
ret += nchoose2(uf.num_val[root]);
}
}
}
return ret + n;
}
};A smarter union-find (AC)
Based on [2] without graph-building optimization (917 ms)
class UF{
public:
vector<int> roots;
UF(int n){
roots.assign(n, 0);for (int i = 0; i < n; ++i)roots[i] = i;
}
int find_set(int x){
return roots[x] == x?x: roots[x] = find_set(roots[x]);
}
void union_set(int x, int y){
roots[find_set(y)] = find_set(x);
}
};
int nChoose2(int n ){
return n*(n-1)/2;
}
class Solution {
public:
unordered_map<int, vector<int>>graph;
map<int, vector<int>>val_nodes;
void build_graph(vector<vector<int>>& edges, vector<int>& vals){
for (auto e:edges){
int u = e[0], v=e[1];
graph[u].push_back(v);
graph[v].push_back(u);
}
}
int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
int ret = 0;
const int n = vals.size();
build_graph(edges, vals);
for (int i = 0 ; i < n; ++i){
val_nodes[vals[i]].push_back(i);
}
UF uf(n);
for (auto [val, nodes]: val_nodes)
for (auto node:nodes){
for (auto nei: graph[node]){
if (vals[nei] > val)continue;
uf.union_set(node, nei);
}
}
unordered_map<int, int> num_vals_in_a_set;
for (auto node: nodes)
num_vals_in_a_set[uf.find_set(node)] ++;
for (auto &[k, v]: num_vals_in_a_set){
ret += nChoose2(v);
}
}
return ret + n;
}
};Based on [2] with graph-building optimization (902 ms)
We can slightly improve the above solution by pruning the neighbor nodes which have a larger value than the current node. It can speed up a little bit and make sure the result is correct. The reason that the results are correct as the tree’s edges are undirected, for the connection from node A to node B, if A’s value is 10 and B’s value is 11. It is fine that there is no path from A to B as there can be a path from B to A, so A and B can still be connected.
class UF{
public:
vector<int> roots;
UF(int n){
roots.assign(n, 0);for (int i = 0; i < n; ++i)roots[i] = i;
}
int find_set(int x){
return roots[x] == x?x: roots[x] = find_set(roots[x]);
}
void union_set(int x, int y){
roots[find_set(y)] = find_set(x);
}
};
int nChoose2(int n ){
return n*(n-1)/2;
}
class Solution {
public:
unordered_map<int, vector<int>>graph;
map<int, vector<int>>val_nodes;
void build_graph(vector<vector<int>>& edges, vector<int>& vals){
for (auto e:edges){
int u = e[0], v=e[1];
if (vals[v] <= vals[u]){
graph[u].push_back(v);
}
if (vals[u] <= vals[v]){
graph[v].push_back(u);
}
}
}
int numberOfGoodPaths(vector<int>& vals, vector<vector<int>>& edges) {
int ret = 0;
const int n = vals.size();
build_graph(edges, vals);
for (int i = 0 ; i < n; ++i){
val_nodes[vals[i]].push_back(i);
}
UF uf(n);
for (auto [val, nodes]: val_nodes){
for (auto node:nodes){
for (auto nei: graph[node]){
uf.union_set(node, nei);
}
}
unordered_map<int, int> num_vals_in_a_set;
for (auto node: nodes)
num_vals_in_a_set[uf.find_set(node)] ++;
for (auto &[k, v]: num_vals_in_a_set){
ret += nChoose2(v);
}
}
return ret + n;
}
};
