#include <cstdio> #include <cstdlib> #include <cassert> #include <algorithm> static const int MAX_N = 100 * 1000; struct vertex { int first_edge; }; struct edge { int to; int next; }; vertex verts[MAX_N]; edge edges[2 * MAX_N]; bool source_colors[MAX_N]; bool dest_colors[MAX_N]; void load_colors(bool* dest, int n) { for (int i = 0; i < n; i++) { int c = getchar(); if (c == '0' || c == '1') { dest[i] = (c == '1'); } else { i--; } } } void load_tree(int n) { for (int i = 0; i < n; i++) { verts[i].first_edge = -1; } for (int i = 0; i < n - 1; i++) { int a, b; assert(2 == scanf("%d %d", &a, &b)); a--; b--; edges[2 * i].to = b; edges[2 * i].next = verts[a].first_edge; verts[a].first_edge = 2 * i; edges[2 * i + 1].to = a; edges[2 * i + 1].next = verts[b].first_edge; verts[b].first_edge = 2 * i + 1; } } bool has_monochrome_edge(const bool* colors, int n) { // We cheat here a bit: we know how edges are allocated, so we can assume // that the edge of id 2*i is a reflection of the edge at position 2*i+1 for (int i = 0; i < n - 1; i++) { const int a = edges[2*i+0].to; const int b = edges[2*i+1].to; if (colors[a] == colors[b]) { return true; } } return false; } // Returns the id of a vertex which starts the path, or -1 if the tree is not a path int is_path(int n) { int num_leaves = 0; int last_leaf = -1; for (int i = 0; i < n; i++) { const int first_edge = verts[i].first_edge; assert(first_edge != -1); const int second_edge = edges[first_edge].next; if (second_edge == -1) { last_leaf = i; num_leaves++; if (num_leaves > 2) { // There are more than 3 leaves, so the tree is not a path return -1; } } } return last_leaf; } bool handle_path(int path_start, int n) { // A segment is a sequence of vertices connected in form of a path, // all vertices having the same color. Count the minimum number of segments // into which the path can be decomposed, both for the source and destination // colorings. auto count_segments = [&] (const bool* colors) { int v = edges[verts[path_start].first_edge].to; bool prev_color = colors[path_start]; int prev_v = path_start; int segment_count = 1; while (true) { if (colors[v] != prev_color) { segment_count++; prev_color = colors[v]; } bool found = false; for (int e = verts[v].first_edge; e != -1; e = edges[e].next) { if (edges[e].to != prev_v) { prev_v = v; v = edges[e].to; found = true; break; } } if (!found) { break; } } return segment_count; }; const int source_count = count_segments(source_colors); const int dest_count = count_segments(dest_colors); // If there are more segments in the destination, we can't solve this. // If there are less, we can always solve this. if (source_count != dest_count) { // printf("Source count: %d, destination count: %d\n", source_count, dest_count); return source_count > dest_count; } // If the number of segments is the same, the only thing that matters // is if the first segment has the same color in both colorings. // printf("Tie: %d vs. %d\n", (int)source_colors[path_start], (int)dest_colors[path_start]); return source_colors[path_start] == dest_colors[path_start]; } bool do_test() { int n; assert(1 == scanf("%d", &n)); load_colors(source_colors, n); load_colors(dest_colors, n); load_tree(n); // Step 1: check if, maybe, the tree is already colored in the right way if (std::equal(source_colors, source_colors + n, dest_colors)) { // No need to do anything, it's already OK // puts("Returning TRUE because the tree is already correctly colored"); return true; } // Step 2: check if the source coloring is monochrome int source_trues = std::count(source_colors, source_colors + n, true); if (source_trues == 0 || source_trues == n) { // The color is monochrome, so we can't change it in any way // puts("Returning FALSE because the source tree is monochrome"); return false; } // Step 3: check if there is a monochrome edge if (!has_monochrome_edge(dest_colors, n)) { // There are no monochrome edges in the target coloring. We know // that we must perform at least one step, and making a step // ensures that there is at least a single monochrome edge, // so it's impossible to achieve the target coloring // puts("Returning FALSE because there is no monochrome edge in the destination coloring"); return false; } // Step 4: check if the tree is a path if (int path_start = is_path(n); path_start != -1) { // The path is a special case return handle_path(path_start, n); } // Otherwise, it is always possible to perform re-coloring. // I have a nice proof of this, but this comment is too short to contain it. // puts("Returning TRUE because the tree is not a path"); return true; } int main() { int t; assert(1 == scanf("%d", &t)); while (t --> 0) { puts(do_test() ? "TAK" : "NIE"); } return 0; }
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 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | #include <cstdio> #include <cstdlib> #include <cassert> #include <algorithm> static const int MAX_N = 100 * 1000; struct vertex { int first_edge; }; struct edge { int to; int next; }; vertex verts[MAX_N]; edge edges[2 * MAX_N]; bool source_colors[MAX_N]; bool dest_colors[MAX_N]; void load_colors(bool* dest, int n) { for (int i = 0; i < n; i++) { int c = getchar(); if (c == '0' || c == '1') { dest[i] = (c == '1'); } else { i--; } } } void load_tree(int n) { for (int i = 0; i < n; i++) { verts[i].first_edge = -1; } for (int i = 0; i < n - 1; i++) { int a, b; assert(2 == scanf("%d %d", &a, &b)); a--; b--; edges[2 * i].to = b; edges[2 * i].next = verts[a].first_edge; verts[a].first_edge = 2 * i; edges[2 * i + 1].to = a; edges[2 * i + 1].next = verts[b].first_edge; verts[b].first_edge = 2 * i + 1; } } bool has_monochrome_edge(const bool* colors, int n) { // We cheat here a bit: we know how edges are allocated, so we can assume // that the edge of id 2*i is a reflection of the edge at position 2*i+1 for (int i = 0; i < n - 1; i++) { const int a = edges[2*i+0].to; const int b = edges[2*i+1].to; if (colors[a] == colors[b]) { return true; } } return false; } // Returns the id of a vertex which starts the path, or -1 if the tree is not a path int is_path(int n) { int num_leaves = 0; int last_leaf = -1; for (int i = 0; i < n; i++) { const int first_edge = verts[i].first_edge; assert(first_edge != -1); const int second_edge = edges[first_edge].next; if (second_edge == -1) { last_leaf = i; num_leaves++; if (num_leaves > 2) { // There are more than 3 leaves, so the tree is not a path return -1; } } } return last_leaf; } bool handle_path(int path_start, int n) { // A segment is a sequence of vertices connected in form of a path, // all vertices having the same color. Count the minimum number of segments // into which the path can be decomposed, both for the source and destination // colorings. auto count_segments = [&] (const bool* colors) { int v = edges[verts[path_start].first_edge].to; bool prev_color = colors[path_start]; int prev_v = path_start; int segment_count = 1; while (true) { if (colors[v] != prev_color) { segment_count++; prev_color = colors[v]; } bool found = false; for (int e = verts[v].first_edge; e != -1; e = edges[e].next) { if (edges[e].to != prev_v) { prev_v = v; v = edges[e].to; found = true; break; } } if (!found) { break; } } return segment_count; }; const int source_count = count_segments(source_colors); const int dest_count = count_segments(dest_colors); // If there are more segments in the destination, we can't solve this. // If there are less, we can always solve this. if (source_count != dest_count) { // printf("Source count: %d, destination count: %d\n", source_count, dest_count); return source_count > dest_count; } // If the number of segments is the same, the only thing that matters // is if the first segment has the same color in both colorings. // printf("Tie: %d vs. %d\n", (int)source_colors[path_start], (int)dest_colors[path_start]); return source_colors[path_start] == dest_colors[path_start]; } bool do_test() { int n; assert(1 == scanf("%d", &n)); load_colors(source_colors, n); load_colors(dest_colors, n); load_tree(n); // Step 1: check if, maybe, the tree is already colored in the right way if (std::equal(source_colors, source_colors + n, dest_colors)) { // No need to do anything, it's already OK // puts("Returning TRUE because the tree is already correctly colored"); return true; } // Step 2: check if the source coloring is monochrome int source_trues = std::count(source_colors, source_colors + n, true); if (source_trues == 0 || source_trues == n) { // The color is monochrome, so we can't change it in any way // puts("Returning FALSE because the source tree is monochrome"); return false; } // Step 3: check if there is a monochrome edge if (!has_monochrome_edge(dest_colors, n)) { // There are no monochrome edges in the target coloring. We know // that we must perform at least one step, and making a step // ensures that there is at least a single monochrome edge, // so it's impossible to achieve the target coloring // puts("Returning FALSE because there is no monochrome edge in the destination coloring"); return false; } // Step 4: check if the tree is a path if (int path_start = is_path(n); path_start != -1) { // The path is a special case return handle_path(path_start, n); } // Otherwise, it is always possible to perform re-coloring. // I have a nice proof of this, but this comment is too short to contain it. // puts("Returning TRUE because the tree is not a path"); return true; } int main() { int t; assert(1 == scanf("%d", &t)); while (t --> 0) { puts(do_test() ? "TAK" : "NIE"); } return 0; } |