#include <bits/stdc++.h>
using namespace std;
#define all(x) x.begin(), x.end()
#define mod 1000000007
#define sp << " " <<
#define endl << '\n'
vector<vector<int>> graph;
vector<long long int> dist;
vector<long long int> distind;
vector<long long int> distmul;
long long int hard = 0;
long long int cnt = 0;
int n;
void dp(int node = 0, int parent = -1){
for (auto go : graph[node]){
if (go == parent)
continue;
dp(go, node);
if (dist[node] < dist[go] + 1){
dist[node] = dist[go] + 1;
distmul[node] = distmul[go];
}else if (dist[node] == dist[go] + 1){
distmul[node] += distmul[go];
}
}
}
void reroot(int node = 0, int parent = -1, long long int parentdist = -1, long long int parentmul = 1){
long long int big1 = parentdist + 1, big2 = 0, big3 = 0;
long long int mul1 = parentmul, mul2 = 0, mul3 = 0;
for (int go : graph[node]){
if (go == parent)
continue;
if (dist[go] + 1 > big1){
big3 = big2;
mul3 = mul2;
big2 = big1;
mul2 = mul1;
big1 = dist[go] + 1;
mul1 = distmul[go];
}else if (dist[go] + 1 > big2){
big3 = big2;
mul3 = mul2;
big2 = dist[go] + 1;
mul2 = distmul[go];
}else if (dist[go] + 1 > big3){
big3 = dist[go] + 1;
mul3 = distmul[go];
}
}
long long int deg1 = big1 * (big2 + big3);
if (graph[node].size() > 2 && deg1 > hard){
hard = deg1;
}
for (int go : graph[node]){
if (go == parent)
continue;
if (big1 == dist[go] + 1){
reroot(go, node, big2, mul2);
}else{
reroot(go, node, big1, mul1);
}
}
}
void reroot2(int node = 0, int parent = -1, long long int parentdist = -1, long long int parentmul = 1, bool parentused = false){
if (graph[node].size() == 0)
return;
long long int big1 = parentdist + 1, big2 = 0, big3 = 0;
long long int mul1 = parentmul, mul2 = 0, mul3 = 0;
long long int smul1 = parentmul, smul2 = 0;
long long int num1 = !parentused, num2 = 0, num3 = 0;
if (parentused){
big1 = 0;
mul1 = 0;
}
for (int go : graph[node]){
if (go == parent)
continue;
if (dist[go] + 1 == big1){
mul1 += distmul[go];
smul1 += distmul[go] * distmul[go];
num1++;
}else if (dist[go] + 1 == big2){
mul2 += distmul[go];
smul2 += distmul[go] * distmul[go];
num2++;
}else if (dist[go] + 1 == big3){
mul3 += distmul[go];
num3++;
}else if (dist[go] + 1 > big1){
big3 = big2;
mul3 = mul2;
num3 = num2;
big2 = big1;
mul2 = mul1;
num2 = num1;
smul2 = smul1;
big1 = dist[go] + 1;
mul1 = distmul[go];
num1 = 1;
smul1 = mul1 * mul1;
}else if (dist[go] + 1 > big2){
big3 = big2;
mul3 = mul2;
num3 = num2;
big2 = dist[go] + 1;
mul2 = distmul[go];
smul2 = mul2 * mul2;
num2 = 1;
}else if (dist[go] + 1 > big3){
big3 = dist[go] + 1;
mul3 = distmul[go];
num3 = 1;
}
}
bool flag1 = false, flag2 = false;
if (parentused && (big1 + big2) * (parentdist + 1) == hard){
cnt += mul2 * mul1;
flag1 = true;
flag2 = true;
}else if (parentused && (num1 > 1 && (2 * big1 * (parentdist + 1)) == hard)){
cnt += (mul1 * mul1 - smul1) / 2;
flag1 = true;
flag2 = true;
}else if (big1 * (big2 + big3) == hard || (num2 > 1 && big1 * big2 * 2 == hard) || (num1 > 1 && big1 * (big1 + big2) == hard) || (num1 > 2 && big1 * big1 * 2 == hard)){
if (num1 > 2){
cnt += (mul1 * mul1 - smul1) / 2;
flag2 = true;
flag1 = true;
}else if (num1 > 1){
cnt += mul2 * mul1;
flag2 = true;
flag1 = true;
}else if (num2 > 1){
cnt += (mul2 * mul2 - smul2) / 2;
flag1 = true;
}else{
cnt += mul2 * mul3;
flag1 = true;
}
}
for (int go : graph[node]){
if (go == parent)
continue;
if (parentused){
if (big1 == dist[go] + 1 && num1 > 1 && big1 >= parentdist)
reroot2(go, node, big1, mul1 - distmul[go], flag2);
else if (big1 != dist[go] + 1 && big1 >= parentdist)
reroot2(go, node, big1, mul1, flag2);
else if (big2 >= parentdist)
reroot2(go, node, big2, mul2, flag1);
else
reroot2(go, node, parentdist + 1, parentmul, true);
}else{
if (big1 == dist[go] + 1 && num1 == 1)
reroot2(go, node, big2, mul2, flag1);
else if (big1 == dist[go] + 1 && num1 > 1)
reroot2(go, node, big1, mul1 - distmul[go], flag2);
else
reroot2(go, node, big1, mul1, flag2);
}
}
}
signed main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
cout.tie(0);
cin >> n;
graph.resize(n);
dist.resize(n, 0);
distind.resize(n);
distmul.resize(n, 1);
for (int i = 0; i < n - 1; i++){
int a, b;
cin >> a >> b;
a--; b--;
graph[a].push_back(b);
graph[b].push_back(a);
}
dp();
reroot();
reroot2();
cout << hard sp cnt;
}
Compilation message
road.cpp: In function 'void reroot(int, int, long long int, long long int)':
road.cpp:39:44: warning: variable 'mul3' set but not used [-Wunused-but-set-variable]
39 | long long int mul1 = parentmul, mul2 = 0, mul3 = 0;
| ^~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
336 KB |
Output is correct |
| 2 |
Correct |
1 ms |
336 KB |
Output is correct |
| 3 |
Incorrect |
1 ms |
336 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
336 KB |
Output is correct |
| 2 |
Correct |
1 ms |
336 KB |
Output is correct |
| 3 |
Incorrect |
1 ms |
336 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
1 ms |
336 KB |
Output is correct |
| 2 |
Correct |
1 ms |
336 KB |
Output is correct |
| 3 |
Incorrect |
1 ms |
336 KB |
Output isn't correct |
| 4 |
Halted |
0 ms |
0 KB |
- |
