Submission #1116242

#	Time	Username	Problem	Language	Result	Execution time	Memory
1116242		hamzabc	Hard route (IZhO17_road)	C++14	100 / 100	643 ms	117364 KiB

road

#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 = 0, bool parentused = true){
	if (graph[node].size() == 0)
		return;
	long long int big1 = parentdist + 1, big2 = 0, big3 = 0;
	long long int mul1 = max(parentmul, 1LL), 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;
		}
	}
	// cerr << node sp "|" sp big1 sp mul1 sp num1 sp "|" sp big2 sp mul2 sp num2 sp "|" sp big3 sp mul3 sp num3 sp "|" sp parentused sp "pd:"sp parentdist + 1 sp parentmul endl;
	bool flag1 = false, flag2 = false;
	if (parentused && (big1 + big2) * (parentdist + 1) == hard){
		cnt += mul2 * mul1;
		if (((parentdist + 1) + big2) * big1 == hard)
			cnt += parentmul * mul2;
		flag1 = true;
		flag2 = true;
	}else if (parentused && (num1 > 1 && (2 * big1 * (parentdist + 1)) == hard)){
		cnt += (mul1 * mul1 - smul1) / 2;
		if (((parentdist + 1) + big2) * big1 == hard)
			cnt += parentmul * mul2;
		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 + 1)
				reroot2(go, node, big1, mul1 - distmul[go], flag2);
			else if (big1 != dist[go] + 1 && big1 >= parentdist + 1){
				reroot2(go, node, big1, mul1, flag2);
			}else if (big2 >= parentdist + 1){
				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 max(1LL, cnt);
}

Compilation message (stderr)

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;
      |                                            ^~~~

• Subtask 1: 19 / 19
• Subtask 2: 33 / 33
• Subtask 3: 48 / 48