#include <bits/stdc++.h> using namespace std; pair<int, int> cnt(string& s) { pair<int, int> res; for (auto& c : s) { if (c == '0') { res.first++; } else { res.second++; } } return res; } int get_deg3plus(vector<vector<int>>& adj) { for (int i = 0; i < adj.size(); i++) { if (adj[i].size() > 2) { return i; } } return -1; } int get_deg1(vector<vector<int>>& adj) { for (int i = 0; i < adj.size(); i++) { if (adj[i].size() == 1) { return i; } } return -1; } vector<int> traverse(vector<vector<int>>& adj, int start) { vector<int> res(adj.size()); res[0] = start; res[1] = adj[start][0]; for (int i = 2; i < adj.size(); i++) { auto& nbrs = adj[res[i - 1]]; res[i] = (nbrs[0] != res[i - 2]) ? nbrs[0] : nbrs[1]; } return res; } bool answer(string& c0, string& c1, vector<vector<int>>& adj) { if (c0 == c1) { return true; } auto p0 = cnt(c0); auto p1 = cnt(c1); if ((p1.first > 0 && p0.first == 0) || (p1.second > 0 && p0.second == 0)) { // missing color needed on final return false; } if (p1.first == 0 || p1.second == 0) { // single color on final return true; } // if (n==1) { // already done... int v3_node = get_deg3plus(adj); // any... if present if (v3_node != -1) { // possible if there exists 2 neighbours with same color for (int i = 0; i < adj.size(); i++) { for (int j : adj[i]) { if (c1[i] == c1[j]) { return true; } } } return false; } // just a line vector<int> line_order = traverse(adj, get_deg1(adj)); // linear check looking for target letters int i0 = 0; for (int i1 = 0; i1 < line_order.size(); i1++) { auto& need_c = c1[line_order[i1]]; while (c0[line_order[i0]] != need_c) { i0++; if (i0 == line_order.size()) { return false; } } } return true; } int main() { // freopen("D:/szkola/testy/dcc/dcc0.in", "r", stdin); ios_base::sync_with_stdio(0); cin.tie(0); int t; cin >> t; while (t--) { int n; cin >> n; string c0, c1; // colors cin >> c0 >> c1; vector<vector<int>> adj(n); for (int i = 0; i < n - 1; i++) { int a, b; cin >> a >> b; a--; b--; adj[a].push_back(b); adj[b].push_back(a); } if (answer(c0, c1, adj)) { 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 | #include <bits/stdc++.h> using namespace std; pair<int, int> cnt(string& s) { pair<int, int> res; for (auto& c : s) { if (c == '0') { res.first++; } else { res.second++; } } return res; } int get_deg3plus(vector<vector<int>>& adj) { for (int i = 0; i < adj.size(); i++) { if (adj[i].size() > 2) { return i; } } return -1; } int get_deg1(vector<vector<int>>& adj) { for (int i = 0; i < adj.size(); i++) { if (adj[i].size() == 1) { return i; } } return -1; } vector<int> traverse(vector<vector<int>>& adj, int start) { vector<int> res(adj.size()); res[0] = start; res[1] = adj[start][0]; for (int i = 2; i < adj.size(); i++) { auto& nbrs = adj[res[i - 1]]; res[i] = (nbrs[0] != res[i - 2]) ? nbrs[0] : nbrs[1]; } return res; } bool answer(string& c0, string& c1, vector<vector<int>>& adj) { if (c0 == c1) { return true; } auto p0 = cnt(c0); auto p1 = cnt(c1); if ((p1.first > 0 && p0.first == 0) || (p1.second > 0 && p0.second == 0)) { // missing color needed on final return false; } if (p1.first == 0 || p1.second == 0) { // single color on final return true; } // if (n==1) { // already done... int v3_node = get_deg3plus(adj); // any... if present if (v3_node != -1) { // possible if there exists 2 neighbours with same color for (int i = 0; i < adj.size(); i++) { for (int j : adj[i]) { if (c1[i] == c1[j]) { return true; } } } return false; } // just a line vector<int> line_order = traverse(adj, get_deg1(adj)); // linear check looking for target letters int i0 = 0; for (int i1 = 0; i1 < line_order.size(); i1++) { auto& need_c = c1[line_order[i1]]; while (c0[line_order[i0]] != need_c) { i0++; if (i0 == line_order.size()) { return false; } } } return true; } int main() { // freopen("D:/szkola/testy/dcc/dcc0.in", "r", stdin); ios_base::sync_with_stdio(0); cin.tie(0); int t; cin >> t; while (t--) { int n; cin >> n; string c0, c1; // colors cin >> c0 >> c1; vector<vector<int>> adj(n); for (int i = 0; i < n - 1; i++) { int a, b; cin >> a >> b; a--; b--; adj[a].push_back(b); adj[b].push_back(a); } if (answer(c0, c1, adj)) { cout << "TAK\n"; } else { cout << "NIE\n"; } } } |