Submission #75925325

---

Source Code Expand

#pragma GCC optimize("Ofast,unroll-loops")
#include <bits/stdc++.h>
#define rep(i, s, e) for (ll i = (s); i <= (e); i++)
#define all(v) v.begin(), v.end()
#define pb push_back
using namespace std;
using ll = long long;
using ull = unsigned long long;
using ld = long double;
using P = array<ll, 2>;
ll n, z, p[101010], u, v, sp[202020][20], h[202020], vs[101010], res, dp[202020], sz[202020];
vector<ll> e[202020], ed[101010];
P dia[101010];
void dfs(ll x, ll f) {
	for (ll i : e[x])
		if (i ^ f)
			h[i] = h[x] + 1, sp[i][0] = x, dfs(i, x);
}
ll dist(P v) {
	auto [x, y] = v;
	ll r = 0;
	if (h[x] < h[y])
		swap(x, y);
	rep(k, 0, 19) if (h[x] - h[y] >> k & 1) x = sp[x][k], r += 1 << k;
	if (x == y)
		return r;
	for (ll k = 20; k--;)
		if (sp[x][k] != sp[y][k])
			x = sp[x][k], y = sp[y][k], r += 2 << k;
	return r + 2;
}
P mg(P a, P b) {
	P r = dist(a) > dist(b) ? a : b;
	for (ll i : a)
		for (ll j : b)
			if (dist({i, j}) > dist(r))
				r = {i, j};
	return r;
}
ll cy(ll x, ll t) {
	if (vs[x])
		return vs[x] == t;
	vs[x] = t;
	for (ll i : ed[x])
		if (cy(i, t))
			return 1;
	return 0;
}
void dc(ll x, ll f) {
	dia[x] = mg(dia[x], dia[f]);
	ll d = dist(dia[x]) / 2;
	auto [u, v] = dia[x];
	res += d;
	if (h[u] < h[v])
		swap(u, v);
	rep(k, 0, 19) if (d >> k & 1) u = sp[u][k];
	sz[u]++;
	for (ll i : ed[x])
		dc(i, x);
}
void gd(ll x, ll f) {
	for (ll i : e[x])
		if (i ^ f) {
			gd(i, x);
			dp[x] += dp[i] + sz[i], sz[x] += sz[i];
		}
}
void rrt(ll x, ll f) {
	for (ll i : e[x])
		if (i ^ f) {
			dp[i] = dp[x] + n - 2 * sz[i], sz[i] = n;
			rrt(i, x);
		}
}
int main() {
	cin.tie(0)->sync_with_stdio(0);
	cin >> n;
	rep(i, 1, n) cin >> p[i], dia[i] = {i, i}, ed[p[i]].push_back(i);
	z = 2 * n - 1;
	rep(i, n + 1, z) {
		cin >> u >> v;
		e[u].pb(i), e[i].pb(u), e[v].pb(i), e[i].pb(v);
	}
	dfs(1, 0);
	rep(k, 1, 19) rep(i, 1, z) sp[i][k] = sp[sp[i][k - 1]][k - 1];
	rep(i, 1, n) if (cy(i, i)) {
		for (ll j = p[i]; j != i; j = p[j])
			dia[i] = mg(dia[i], dia[j]);
		ed[p[i]].erase(ranges::find(ed[p[i]], i));
		for (ll j = p[i]; j != i; j = p[j]) {
			ed[p[j]].erase(ranges::find(ed[p[j]], i));
		}
		dc(i, i);
		for (ll j = p[i]; j != i; j = p[j])
			dc(j, i);
	}
	gd(1, 0);
	rrt(1, 0);
	rep(i, 1, n) cout << (res + dp[i]) / 2 << ' ';
	return 0;
}

Submission Info

Compile Error

./Main.cpp: In function 'll dist(P)':
./Main.cpp:24:32: warning: suggest parentheses around '-' inside '>>' [-Wparentheses]
   24 |         rep(k, 0, 19) if (h[x] - h[y] >> k & 1) x = sp[x][k], r += 1 << k;
      |                           ~~~~~^~~~~~

Judge Result