Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #include <algorithm>
- #include <iostream>
- #include <vector>
- using namespace std;
- const int N = 222222;
- int n, f[N], p[N], d;
- long long result;
- vector<int> g[N];
- void input() {
- cin >> n;
- for (int i = 1; i < n; ++i) {
- int x, y;
- cin >> x >> y;
- g[x].push_back(y);
- g[y].push_back(x);
- }
- }
- void calc_fp(int v=1, int parent=-1) {
- p[v] = 1;
- for (size_t i = 0; i < g[v].size(); ++i) {
- int to = g[v][i];
- if (to != parent) {
- calc_fp(to, v);
- if (f[to] + 1 == f[v]) {
- p[v] += p[to];
- } else if (f[to] + 1 > f[v]) {
- f[v] = f[to] + 1;
- p[v] = p[to];
- }
- }
- }
- }
- void calc_diam(int v=1, int parent=-1) {
- int max1 = 0, max2 = 0;
- for (size_t i = 0; i < g[v].size(); ++i) {
- int to = g[v][i];
- if (to != parent) {
- calc_diam(to, v);
- max2 = max(max2, f[to] + 1);
- if (max2 > max1) {
- swap(max1, max2);
- }
- }
- }
- d = max(d, max1 + max2);
- }
- void calc_result(int v=1, int parent=-1) {
- int max1 = 0, max2 = 0, p_sum = 0;
- for (size_t i = 0; i < g[v].size(); ++i) {
- int to = g[v][i];
- if (to != parent) {
- calc_result(to, v);
- p_sum += p[to];
- max2 = max(max2, f[to] + 1);
- if (max2 > max1) {
- swap(max1, max2);
- }
- }
- }
- if (max1 + max2 == d) {
- if (max2 == 0) {
- result += p_sum;
- } else if (max1 != max2) {
- p_sum = 0;
- int p_max = 0;
- for (size_t i = 0; i < g[v].size(); ++i) {
- int to = g[v][i];
- if (to != parent) {
- if (f[to] + 1 == max2) {
- p_sum += p[to];
- } else if (f[to] + 1 == max1) {
- p_max = p[to];
- }
- }
- }
- result += p_max * 1ll * p_sum;
- } else if (max1 == max2) {
- p_sum = 0;
- long long temp = 0;
- for (size_t i = 0; i < g[v].size(); ++i) {
- int to = g[v][i];
- if (to != parent && f[to] + 1 == max1) {
- p_sum += p[to];
- }
- }
- for (size_t i = 0; i < g[v].size(); ++i) {
- int to = g[v][i];
- if (to != parent && f[to] + 1 == max1) {
- temp += p[to] * 1ll * (p_sum - p[to]);
- }
- }
- result += temp / 2;
- }
- }
- }
- int main() {
- input();
- calc_fp();
- calc_diam();
- calc_result();
- cout << result << endl;
- return 0;
- }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement