wannafly Day2 E 阔力梯的树

dsu on tree

题目链接

点我跳转

题目大意

解题思路

AC_Code

#include<bits/stdc++.h>
#define rep(i,a,n) for (int i=a;i<=n;i++)
#define int long long
using namespace std;
const int N = 3e5 + 10 , inf = 0x3f3f3f3f;
struct Edge{
int nex , to;
}edge[N << 1];
int head
 , TOT;
void add_edge(int u , int v)
{
edge[++ TOT].nex = head[u] ;
edge[TOT].to = v;
head[u] = TOT;
}
int hson
 , dep
 , sz
 , HH;
int n , ans
 , now;
set<int>se;
void dfs(int u , int far)
{
dep[u] = dep[far] + 1;
sz[u] = 1;
for(int i = head[u] ; i ; i = edge[i].nex)
{
int v = edge[i].to;
if(v == far) continue ;
dfs(v , u);
sz[u] += sz[v];
if(sz[v] > sz[hson[u]]) hson[u] = v;
}
}
void calc(int u , int far , int val)
{
if(val == 1)
{
auto it = se.upper_bound(u);
auto it1 = it , it2 = it;
it1 -- ;
int add = 0 , add1 = 0 , add2 = 0;
if(it1 != se.begin() && it2 != se.end()) add = *it2 - *it1;
now -= add * add;
if(it1 != se.begin()) add1 = u - *it1;
if(it2 != se.end())   add2 = *it2 - u;
now += add1 * add1 + add2 * add2;
se.insert(u);
}
else if(val == -1)
{
se.erase(u);
auto it = se.upper_bound(u);
auto it1 = it , it2 = it;
it1 -- ;
int add = 0 , add1 = 0 , add2 = 0;
if(it1 != se.begin()) add1 = u - *it1;
if(it2 != se.end())   add2 = *it2 - u;
now -= add1 * add1 , now -= add2 * add2;
if(it1 != se.begin() && it2 != se.end()) add = *it2 - *it1;
now += add * add;
}
for(int i = head[u] ; i ; i = edge[i].nex)
{
int v = edge[i].to;
if(v == far || v == HH) continue ;
calc(v , u , val);
}
}
void dsu(int u , int far , int op)
{
for(int i = head[u] ; i ; i = edge[i].nex)
{
int v = edge[i].to;
if(v == far || v == hson[u]) continue ;
dsu(v , u , 0);
}
if(hson[u]) dsu(hson[u] , u , 1) , HH = hson[u];
calc(u , far , 1);
ans[u] = now;
HH = 0;
if(!op) calc(u , far , -1) , now = 0;
}
signed main()
{
cin >> n;
se.insert(-inf);
rep(i , 2 , n)
{
int x;
cin >> x;
add_edge(i , x) , add_edge(x , i);
}
dfs(1 , 0);
dsu(1 , 0 , 0);
rep(i , 1 , n) cout << ans[i] << '\n';
return 0;
}