Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <bits/stdc++.h>
- long long solveYuri(int64_t n, int x, int y, std::vector<std::pair<int, int>> arr) {
- std::vector<std::vector<std::pair<int, bool>>> graph(300000);
- int64_t sum, a, b, c, d;
- // Удаляет путь от a в b
- std::function<bool(int, int64_t)> find = [&](int from , int64_t to) {
- int i=0;
- if(to == b) {
- for(; graph[to][i].first != from; i++);
- graph[to][i].second = false;
- return false;
- }
- bool fnd = true;
- for(; i<graph[to].size() && fnd; i++) {
- if (graph[to][i].first != from) {
- fnd = find(to, graph[to][i].first);
- if (!fnd) {
- graph[to][i].second = false;
- if (from != -1) {
- int j=0;
- for(; graph[to][j].first != from; j++);
- graph[to][j].second = false;
- }
- // std::cout << to << ' ' << graph[to][i].first << '\n';
- }
- }
- }
- return fnd;
- };
- // Сколько вершин в компоненте связности
- std::function<int64_t(int, int64_t)> many = [&](int from , int64_t to) {
- int i=0;
- if (from != -1) {
- for(; graph[to][i].first != from; i++);
- graph[to][i].second = false;
- }
- int64_t s = 1;
- for(i=0; i<graph[to].size() && graph[to][i].first != from; i++) {
- if (graph[to][i].second) {
- s += many(to, graph[to][i].first);
- graph[to][i].second = false;
- // std::cout << to << ' ' << graph[to][i].first << '\n';
- }
- }
- return s;
- };
- a = x-1;
- b = y-1;
- sum = n * (n-1);
- for (auto& p : arr) {
- c = p.first;
- d = p.second;
- c--;
- d--;
- graph[c].push_back({d, true});
- graph[d].push_back({c, true});
- }
- find(-1, a);
- int64_t s1 = many(-1, a);
- int64_t s2 = many(-1, b);
- return sum - s1 * s2;
- }
- long long solveDmitry(int n, int x, int y, const std::vector<std::pair<int, int>>& arr) {
- --x, --y;
- std::vector<std::vector<int>> edges(n);
- for (int i = 0; i < n-1; ++i) {
- int a = arr[i].first-1, b = arr[i].second-1;
- edges[a].push_back(b);
- edges[b].push_back(a);
- }
- std::vector<int> size(n); // size[v] = size subtree of vertex v
- std::vector<bool> is_x(n); // is_x[v] = true if vertex x in subtree of vertex v
- std::function<void(int, int)> visit = [&](const int curr, const int parent) {
- size[curr]++;
- if (curr == x) {
- is_x[curr] = true;
- }
- for (auto next : edges[curr]) {
- if (next != parent) {
- visit(next, curr);
- size[curr] += size[next];
- if (is_x[next] && !is_x[curr]) { // vertex x in subtree[curr] because vertex x in subtree[next]
- is_x[curr] = is_x[next];
- }
- }
- }
- };
- visit(y, y);
- long long sum_without_x = 1;
- long long count_with_x = 0; // for check myself
- for (auto next : edges[y]) {
- if (!is_x[next]) {
- sum_without_x += size[next];
- } else {
- count_with_x++;
- }
- }
- assert(count_with_x == 1); // check myself
- // can't wolk from subtree x in subtree[y] \ subtree[v], which is_x[v] = false
- return 1LL * n * (n-1) - 1LL * size[x] * sum_without_x;
- }
- /*
- void gen(int n, int& x, int& y,std::vector<std::pair<int, int>>& arr) {
- arr.assign(
- }
- */
- int main() {
- int n, x, y;
- std::vector<std::pair<int, int>> arr;
- std::cin >> n >> x >> y;
- for (int i = 0; i < n-1; ++i) {
- int a, b; std::cin >> a >> b;
- arr.push_back({a, b});
- }
- std::cout << "Dmitry answer = " << solveDmitry(n, x, y, arr) << ", but Yuri answer = " << solveYuri(n, x, y, arr) << std::endl;
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement