#include "testlib.h" #include #include #include #include typedef unsigned long long cost_t; struct cost_line { unsigned v1, v2, cost; }; bool vec_contains(std::vector const &vec, unsigned val) { for (unsigned e : vec) { if (e == val) { return true; } } return false; } unsigned erase(std::vector &vec, unsigned val) { auto it = std::remove(vec.begin(), vec.end(), val); unsigned rv = vec.end() - it; vec.erase(it, vec.end()); return rv; } struct graph { std::vector< std::vector > adj; std::vector outflow; std::vector color; graph(unsigned n) : adj(n), outflow(n), color(n) {} bool can_add_edge(unsigned v1, unsigned v2) const { return v1 != v2 && adj[v1].size() < 3 && adj[v2].size() < 3 && !vec_contains(adj[v1], v2) && !vec_contains(adj[v2], v1); } void add_edge(unsigned v1, unsigned v2) { if (!can_add_edge(v1, v2)) { quitf(_fail, "Internal error: adding edge (%u, %u) disallowed", v1 + 1, v2 + 1); } adj[v1].push_back(v2); adj[v2].push_back(v1); } void remove_edge(unsigned v1, unsigned v2) { if (erase(adj[v1], v2) != 1 || erase(adj[v2], v1) != 1) { quitf(_fail, "Internal error: removing edge (%u, %u) failed", v1 + 1, v2 + 1); } } void dfs_color(unsigned v, unsigned c) { color[v] = c; for (unsigned next : adj[v]) { if (color[next] == 0) { dfs_color(next, c); } } } void recolor() { for (unsigned i = 0; i < color.size(); ++i) { color[i] = 0; } unsigned col = 0; for (unsigned i = 0; i < color.size(); ++i) { if (color[i] == 0) { dfs_color(i, ++col); } } if (col != 2) { quitf(_fail, "Internal error: coloring failed, num of colors %u", col); } } void fill_costs(std::vector const &costs) { for (unsigned i = 0; i < outflow.size(); ++i) { outflow[i] = 0; } for (unsigned i = 0; i < costs.size(); ++i) { if (color[costs[i].v1] != color[costs[i].v2]) { outflow[costs[i].v1] += costs[i].cost; outflow[costs[i].v2] += costs[i].cost; } } } bool can_add_answer_edge(unsigned v1, unsigned v2) const { return can_add_edge(v1, v2) && color[v1] != color[v2]; } cost_t dfs_cost(unsigned v, unsigned p, cost_t &aggr) const { cost_t result = outflow[v]; for (unsigned next : adj[v]) { if (next != p) { cost_t rec = dfs_cost(next, v, aggr); aggr += rec; result += rec; } } return result; } cost_t cost_of_new_edge(unsigned v1, unsigned v2) const { if (!can_add_answer_edge(v1, v2)) { quitf(_fail, "Internal error: cannot add answer edge (%u, %u)", v1 + 1, v2 + 1); } cost_t aggr = 0; dfs_cost(v1, v2, aggr); dfs_cost(v2, v1, aggr); return aggr; } }; int main(int argc, char *argv[]) { setName("Checker for 'Network Upgrade'"); registerTestlibCmd(argc, argv); unsigned n = inf.readInt(); graph g(n); std::vector v1s(n - 1), v2s(n - 1); for (unsigned i = 1; i < n; ++i) { unsigned v1 = inf.readInt() - 1; unsigned v2 = inf.readInt() - 1; g.add_edge(v1, v2); v1s[i - 1] = v1; v2s[i - 1] = v2; } unsigned d = inf.readInt(); std::vector costs(d); for (unsigned i = 0; i < d; ++i) { costs[i].v1 = inf.readInt() - 1; costs[i].v2 = inf.readInt() - 1; costs[i].cost = inf.readInt(); } for (unsigned i = 1; i < n; ++i) { unsigned jury_v1 = ans.readInt(1, n) - 1; unsigned jury_v2 = ans.readInt(jury_v1 + 2, n) - 1; unsigned cont_v1 = ouf.readInt(1, n) - 1; unsigned cont_v2 = ouf.readInt(cont_v1 + 2, n) - 1; g.remove_edge(v1s[i - 1], v2s[i - 1]); g.recolor(); if (!g.can_add_answer_edge(jury_v1, jury_v2)) { quitf(_fail, "Jury error: cannot add answer edge (%u, %u) in line %u", jury_v1 + 1, jury_v2 + 1, i); } if (!g.can_add_answer_edge(cont_v1, cont_v2)) { quitf(_wa, "Wrong answer: cannot add answer edge (%u, %u) in line %u", cont_v1 + 1, cont_v2 + 1, i); } // the remaining checks are about costs, so if the answer is the same, // we can skip them if (jury_v1 != cont_v1 || jury_v2 != cont_v2) { g.fill_costs(costs); cost_t jury_cost = g.cost_of_new_edge(jury_v1, jury_v2); cost_t cont_cost = g.cost_of_new_edge(cont_v1, cont_v2); if (cont_cost < jury_cost) { quitf(_fail, "Contestant found a better answer (%u, %u) with cost %llu than jury (%u, %u) with cost %llu", cont_v1 + 1, cont_v2 + 1, cont_cost, jury_v1 + 1, jury_v2 + 1, jury_cost); } if (cont_cost > jury_cost) { quitf(_wa, "Jury has a better answer (%u, %u) with cost %llu than contestant (%u, %u) with cost %llu", jury_v1 + 1, jury_v2 + 1, jury_cost, cont_v1 + 1, cont_v2 + 1, cont_cost); } if (jury_v1 > cont_v1) { quitf(_fail, "Contestant found a lex smaller answer (%u, %u) than jury (%u, %u)", cont_v1 + 1, cont_v2 + 1, jury_v1 + 1, jury_v2 + 1); } if (jury_v1 < cont_v1) { quitf(_wa, "Jury has a lex smaller answer (%u, %u) than contestant (%u, %u)", jury_v1 + 1, jury_v2 + 1, cont_v1 + 1, cont_v2 + 1); } if (jury_v2 > cont_v2) { quitf(_fail, "Contestant found a lex smaller answer (%u, %u) than jury (%u, %u)", cont_v1 + 1, cont_v2 + 1, jury_v1 + 1, jury_v2 + 1); } if (jury_v2 < cont_v2) { quitf(_wa, "Jury has a lex smaller answer (%u, %u) than contestant (%u, %u)", jury_v1 + 1, jury_v2 + 1, cont_v1 + 1, cont_v2 + 1); } } // finally, add the new edge g.add_edge(jury_v1, jury_v2); } quitf(_ok, "All replacements match"); }