// --- algolib/union-find.h ---
// union find, optionally with tags for each component. this is useful
// for example when compressing components. enable using template parameter
#include <algorithm>
#include <type_traits>
#include <numeric>
#include <vector>
template<typename T = void, typename enable = void>
class UnionFind {
std::vector<int> p, size_;
public:
UnionFind(int n = 0) : p(n), size_(n, 1) {
std::iota(p.begin(), p.end(), 0);
}
int create() {
int r = p.size();
p.push_back(r);
size_.push_back(1);
return r;
}
int find(int i) {
if (i == p[i])
return i;
return p[i] = find(p[i]);
}
int operator[](int i) { return find(i); }
int size(int i) { return size_[find(i)]; }
bool connected(int a, int b) { return find(a) == find(b); }
bool connect(int a, int b) {
a = find(a), b = find(b);
if (a == b)
return false;
if (size_[a] > size_[b])
std::swap(a, b);
size_[b] += size_[a];
p[a] = b;
return true;
}
};
template<typename T>
class UnionFind<T, typename std::enable_if_t<!std::is_same_v<T, void>>>
: public UnionFind<void> {
std::vector<T> tags;
public:
UnionFind(int n = 0) : UnionFind<void>(n), tags(n) { }
void set_tag(int i, const T& t) { tags[find(i)] = t; }
const T& tag(int i) { return tags[find(i)]; }
};
// ----------------
#include <bits/stdc++.h>
using namespace std;
const int MAX_N = 300005;
int n;
vector<int> ans;
struct Edge {
int to, key; // key required
};
vector<Edge> out[MAX_N]; // outgoing edges. TODO optimise
int vis[MAX_N], par[MAX_N]; // vis_time, parent in DAG/tree
set<int> keys[MAX_N];
std::vector<int> in_comp[MAX_N]; // for each vertex the set of vertices
// it represents
// weakly is weakly connected components
UnionFind<int> compressed(MAX_N), weakly(MAX_N);
// merge i into j
// TODO smaller-to-larger
void merge(int i, int j) {
compressed.connect(i, j);
compressed.set_tag(j, j);
in_comp[j].insert(in_comp[j].end(),
in_comp[i].begin(), in_comp[i].end());
in_comp[i].clear();
for (int x : keys[i])
keys[j].insert(x);
keys[i].clear();
// TODO no duplicates
out[j].insert(out[j].end(), out[i].begin(), out[i].end());
out[i].clear();
}
void solve() {
bool trivial = false;
for (int i = 0; i < n; i++) {
bool found_out = false;
for (auto& e : out[i])
if (e.key == *keys[i].begin()) { // keys has size 1
found_out = true;
par[i] = e.to;
weakly.connect(i, e.to);
break;
}
if (!found_out) {
trivial = true;
ans[i] = 1;
}
}
if (trivial)
return;
// otherwise each weakly connected component is a tree
// with one extra edge. compress corresponding cycle to get a
// directed forest
// tags are the roots (where keys & edges are stored)
for (int i = 0; i < n; i++) {
compressed.set_tag(i, i);
in_comp[i].push_back(i);
}
// roots, ordered by size
struct PQKey {
int vertex, size;
bool operator<(const PQKey& o) const { return size > o.size; }
};
priority_queue<PQKey> roots;
for (int i = 0; i < n; i++)
if (vis[i] == 0) {
int t = i + 1; // 1-based vis times
int j = i;
do {
vis[j] = t;
j = par[j];
} while (vis[j] == 0);
if (vis[j] == t) { // cycle!
int start = j;
for (j = par[j]; j != start; j = par[j])
merge(j, start);
par[start] = -1;
roots.push({ start, compressed.size(start) });
}
}
// != -1 means we already found an answer of size done_size
// so stop pushing to pq, but collect remaining components
int done_size = -1;
while (!roots.empty()) {
auto [r, _] = roots.top();
roots.pop();
// find any outgoing edge
bool has_outgoing = false;
for (auto& e : out[r]) {
// TODO skip useless edges/duplicates
if (compressed.connected(r, e.to) ||
keys[r].find(e.key) == keys[r].end())
continue;
if (weakly.connected(r, e.to)) {
// merge with self
// walk up parent chain, until we're in root comp
int j = compressed.tag(e.to);
do {
int jp = compressed.tag(par[j]);
merge(j, jp);
j = jp;
} while (par[j] != -1);
// we're still a root, but of larger size
roots.push({ r, compressed.size(r) });
} else { // other component
has_outgoing = true;
// merge with other component, if not done
if (done_size != -1)
break;
has_outgoing = true;
par[r] = e.to;
weakly.connect(r, e.to);
// nothing new for priority queue, as we're
// not a root anymore
break;
}
}
if (!has_outgoing && (done_size == -1 ||
done_size == compressed.size(r))) {
done_size = compressed.size(r);
for (int x : in_comp[r])
ans[x] = 1;
}
}
}
vector<int> find_reachable(vector<int> r, vector<int> u,
vector<int> v, vector<int> c) {
n = r.size();
for (int i = 0; i < n; i++)
keys[i].insert(r[i]);
int m = u.size();
for (int i = 0; i < m; i++) {
out[u[i]].push_back({ v[i], c[i] });
out[v[i]].push_back({ u[i], c[i] });
}
ans.assign(n, 0);
solve();
return ans;
}
