#include <bits/stdc++.h>
using namespace std;
vector<int> c, rep, cnt, is_col_done;
vector<unordered_map<int, int>> col_rep;
vector<vector<int>> pol;
vector<vector<int>> col_list;
int find(int x) {
if (x != rep[x]) {
rep[x] = find(rep[x]);
}
return rep[x];
}
void uni(int a, int b, auto& done) {
a = find(a);
b = find(b);
if (a == b) {
return;
}
if (cnt[a] < cnt[b]) swap(a, b);
rep[b] = a;
cnt[a] += cnt[b];
if (col_rep[a].size() < col_rep[b].size()) swap(col_rep[a], col_rep[b]);
for (auto [i, j] : col_rep[b]) {
auto it = col_rep[a].find(i);
if (it != col_rep[a].end()) {
uni(j, it->second, done);
if (cnt[find(j)] == (int)col_list[c[j]].size() && !is_col_done[c[j]]) {
done.push_back(j);
is_col_done[c[j]] = 1;
}
}
col_rep[a][i] = find(j);
}
}
int main()
{
ios_base::sync_with_stdio(false);
cin.tie(NULL);
int t; cin >> t;
while (t--) {
int n, m, k; cin >> n >> m >> k;
int res = n;
c = vector<int>(n + 1);
rep = vector<int>(n + 1);
cnt = vector<int>(n + 1);
col_rep = vector<unordered_map<int, int>>(n + 1);
pol = vector(n + 1, vector<int>());
col_list = vector(k + 1, vector<int>());
for (int i = 1; i <= n; i++) {
cin >> c[i];
col_list[c[i]].push_back(i);
rep[i] = i;
cnt[i] = 1;
}
for (int i = 0; i < m; i++) {
int a, b; cin >> a >> b;
pol[a].push_back(b);
pol[b].push_back(a);
vector<int> trash;
if (c[a] == c[b]) {
uni(a, b, trash);
}
}
vector<int> done;
is_col_done = vector<int>(k + 1);
for (int i = 1; i <= k; i++) {
// ciekawe czy testy z forum to wykrywaja
// cout << i << " " << col_list[i].size() << endl;
if (col_list[i].size()) {
int sz = cnt[find(col_list[i][0])];
// cout << i << " " << col_list[i].size() << " " << sz << endl;
if (sz == (int)col_list[i].size()) {
done.push_back(col_list[i][0]);
is_col_done[i] = 1;
}
}
}
while (!done.empty()) {
int w = done.back();
// cout << "DONE " << w << " " << c[w] << endl;
// for (auto [i, j] : col_rep[c[w]]) {
// cout << " col_rep " << c[w] << " " << i << " " << j << endl;
// }
done.pop_back();
res -= (int)col_list[c[w]].size();
for (auto i : col_list[c[w]]) {
for (auto j : pol[i]) {
// cout << " " << i << " " << j << " " << (col_rep[find(w)].find(c[j]) != col_rep[find(w)].end() ? col_rep[find(w)][c[j]] : 0) << endl;
if (!col_rep[find(w)][c[j]]) {
col_rep[find(w)][c[j]] = j;
}
else {
uni(j, col_rep[find(w)][c[j]], done);
}
// cout << " " << i << " " << j << " " << (col_rep[find(w)].find(c[j]) != col_rep[find(w)].end() ? col_rep[find(w)][c[j]] : 0) << endl;
}
}
for (auto i : col_list[c[w]]) {
for (auto j : pol[i]) {
// cout << " " << i << " " << j << " " << (col_rep[find(w)].find(c[j]) != col_rep[find(w)].end() ? col_rep[find(w)][c[j]] : 0) << endl;
if (is_col_done[c[j]]) {
uni(i, j, done);
}
else {
int sz = cnt[find(j)];
int sz_total = (int)col_list[c[j]].size();
if (sz == sz_total) {
done.push_back(j);
is_col_done[c[j]] = 1;
}
}
}
}
// for (auto [i, j] : col_rep[c[w]]) {
// cout << " col_rep " << c[w] << " " << i << " " << j << endl;
// }
}
if (res == 0) cout << "TAK\n";
else cout << "NIE\n";
}
}
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | #include <bits/stdc++.h> using namespace std; vector<int> c, rep, cnt, is_col_done; vector<unordered_map<int, int>> col_rep; vector<vector<int>> pol; vector<vector<int>> col_list; int find(int x) { if (x != rep[x]) { rep[x] = find(rep[x]); } return rep[x]; } void uni(int a, int b, auto& done) { a = find(a); b = find(b); if (a == b) { return; } if (cnt[a] < cnt[b]) swap(a, b); rep[b] = a; cnt[a] += cnt[b]; if (col_rep[a].size() < col_rep[b].size()) swap(col_rep[a], col_rep[b]); for (auto [i, j] : col_rep[b]) { auto it = col_rep[a].find(i); if (it != col_rep[a].end()) { uni(j, it->second, done); if (cnt[find(j)] == (int)col_list[c[j]].size() && !is_col_done[c[j]]) { done.push_back(j); is_col_done[c[j]] = 1; } } col_rep[a][i] = find(j); } } int main() { ios_base::sync_with_stdio(false); cin.tie(NULL); int t; cin >> t; while (t--) { int n, m, k; cin >> n >> m >> k; int res = n; c = vector<int>(n + 1); rep = vector<int>(n + 1); cnt = vector<int>(n + 1); col_rep = vector<unordered_map<int, int>>(n + 1); pol = vector(n + 1, vector<int>()); col_list = vector(k + 1, vector<int>()); for (int i = 1; i <= n; i++) { cin >> c[i]; col_list[c[i]].push_back(i); rep[i] = i; cnt[i] = 1; } for (int i = 0; i < m; i++) { int a, b; cin >> a >> b; pol[a].push_back(b); pol[b].push_back(a); vector<int> trash; if (c[a] == c[b]) { uni(a, b, trash); } } vector<int> done; is_col_done = vector<int>(k + 1); for (int i = 1; i <= k; i++) { // ciekawe czy testy z forum to wykrywaja // cout << i << " " << col_list[i].size() << endl; if (col_list[i].size()) { int sz = cnt[find(col_list[i][0])]; // cout << i << " " << col_list[i].size() << " " << sz << endl; if (sz == (int)col_list[i].size()) { done.push_back(col_list[i][0]); is_col_done[i] = 1; } } } while (!done.empty()) { int w = done.back(); // cout << "DONE " << w << " " << c[w] << endl; // for (auto [i, j] : col_rep[c[w]]) { // cout << " col_rep " << c[w] << " " << i << " " << j << endl; // } done.pop_back(); res -= (int)col_list[c[w]].size(); for (auto i : col_list[c[w]]) { for (auto j : pol[i]) { // cout << " " << i << " " << j << " " << (col_rep[find(w)].find(c[j]) != col_rep[find(w)].end() ? col_rep[find(w)][c[j]] : 0) << endl; if (!col_rep[find(w)][c[j]]) { col_rep[find(w)][c[j]] = j; } else { uni(j, col_rep[find(w)][c[j]], done); } // cout << " " << i << " " << j << " " << (col_rep[find(w)].find(c[j]) != col_rep[find(w)].end() ? col_rep[find(w)][c[j]] : 0) << endl; } } for (auto i : col_list[c[w]]) { for (auto j : pol[i]) { // cout << " " << i << " " << j << " " << (col_rep[find(w)].find(c[j]) != col_rep[find(w)].end() ? col_rep[find(w)][c[j]] : 0) << endl; if (is_col_done[c[j]]) { uni(i, j, done); } else { int sz = cnt[find(j)]; int sz_total = (int)col_list[c[j]].size(); if (sz == sz_total) { done.push_back(j); is_col_done[c[j]] = 1; } } } } // for (auto [i, j] : col_rep[c[w]]) { // cout << " col_rep " << c[w] << " " << i << " " << j << endl; // } } if (res == 0) cout << "TAK\n"; else cout << "NIE\n"; } } |
English