#include <iostream> #include <vector> using namespace std; using Tree = vector<vector<int>>; string scolors, ecolors; bool is_line(const Tree &tree) { int leafs = 0; for (const auto &n : tree) { if (size(n) == 1) { ++leafs; if (leafs > 2) return false; } } return true; } int leaf(const Tree &tree) { for (int i = 1; i < size(tree); ++i) { if (size(tree[i]) == 1) return i; } return -1; } int groups(const Tree &tree, const string &colors, int leaf) { int next = tree[leaf][0]; int ret = 1; while (true) { if (colors[leaf] != colors[next]) ++ret; if (size(tree[next]) == 1) break; int tmp = next; next = tree[next][tree[next][0] == leaf]; leaf = tmp; } return ret; } void handle_line(const Tree &tree) { int curr = leaf(tree); char scolor = scolors[curr], ecolor = ecolors[curr]; int sgroups = groups(tree, scolors, curr), egroups = groups(tree, ecolors, curr); cout << (sgroups > egroups || sgroups == egroups && scolor == ecolor ? "TAK" : "NIE") << '\n'; } void handle_non_line(const Tree &tree) { bool ssingle_color = true; for (int i = 2; i < size(scolors); ++i) { if (scolors[i] != scolors[1]) { ssingle_color = false; break; } } bool esingle_color = true; for (int i = 2; i < size(ecolors); ++i) { if (ecolors[i] != ecolors[1]) { esingle_color = false; break; } } if (ssingle_color) { cout << (esingle_color && scolors[1] == ecolors[1] ? "TAK" : "NIE") << '\n'; return; } if (esingle_color) { cout << "TAK" << '\n'; return; } for (int i = 1; i < size(tree); ++i) { for (int j = 0; j < size(tree[i]); ++j) { if (ecolors[i] == ecolors[tree[i][j]]) { cout << "TAK" << '\n'; return; } } } cout << (scolors == ecolors ? "TAK" : "NIE") << '\n'; } void test() { int n; cin >> n; scolors.resize(n + 1); ecolors.resize(n + 1); for (int i = 1; i <= n; ++i) { cin >> scolors[i]; } for (int i = 1; i <= n; ++i) { cin >> ecolors[i]; } Tree tree; tree.resize(n + 1); for (int i = 0; i < n - 1; ++i) { int f, t; cin >> f >> t; tree[f].push_back(t); tree[t].push_back(f); } if (n == 1) { cout << (scolors == ecolors ? "TAK" : "NIE") << '\n'; } else if (is_line(tree)) { handle_line(tree); } else { handle_non_line(tree); } } int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); int t; cin >> t; for (int i = 0; i < t; ++i) { test(); } }
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 | #include <iostream> #include <vector> using namespace std; using Tree = vector<vector<int>>; string scolors, ecolors; bool is_line(const Tree &tree) { int leafs = 0; for (const auto &n : tree) { if (size(n) == 1) { ++leafs; if (leafs > 2) return false; } } return true; } int leaf(const Tree &tree) { for (int i = 1; i < size(tree); ++i) { if (size(tree[i]) == 1) return i; } return -1; } int groups(const Tree &tree, const string &colors, int leaf) { int next = tree[leaf][0]; int ret = 1; while (true) { if (colors[leaf] != colors[next]) ++ret; if (size(tree[next]) == 1) break; int tmp = next; next = tree[next][tree[next][0] == leaf]; leaf = tmp; } return ret; } void handle_line(const Tree &tree) { int curr = leaf(tree); char scolor = scolors[curr], ecolor = ecolors[curr]; int sgroups = groups(tree, scolors, curr), egroups = groups(tree, ecolors, curr); cout << (sgroups > egroups || sgroups == egroups && scolor == ecolor ? "TAK" : "NIE") << '\n'; } void handle_non_line(const Tree &tree) { bool ssingle_color = true; for (int i = 2; i < size(scolors); ++i) { if (scolors[i] != scolors[1]) { ssingle_color = false; break; } } bool esingle_color = true; for (int i = 2; i < size(ecolors); ++i) { if (ecolors[i] != ecolors[1]) { esingle_color = false; break; } } if (ssingle_color) { cout << (esingle_color && scolors[1] == ecolors[1] ? "TAK" : "NIE") << '\n'; return; } if (esingle_color) { cout << "TAK" << '\n'; return; } for (int i = 1; i < size(tree); ++i) { for (int j = 0; j < size(tree[i]); ++j) { if (ecolors[i] == ecolors[tree[i][j]]) { cout << "TAK" << '\n'; return; } } } cout << (scolors == ecolors ? "TAK" : "NIE") << '\n'; } void test() { int n; cin >> n; scolors.resize(n + 1); ecolors.resize(n + 1); for (int i = 1; i <= n; ++i) { cin >> scolors[i]; } for (int i = 1; i <= n; ++i) { cin >> ecolors[i]; } Tree tree; tree.resize(n + 1); for (int i = 0; i < n - 1; ++i) { int f, t; cin >> f >> t; tree[f].push_back(t); tree[t].push_back(f); } if (n == 1) { cout << (scolors == ecolors ? "TAK" : "NIE") << '\n'; } else if (is_line(tree)) { handle_line(tree); } else { handle_non_line(tree); } } int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); int t; cin >> t; for (int i = 0; i < t; ++i) { test(); } } |