2020 CCPC Wannafly Winter Camp Day2 E阔力梯的树（树上启发式合并）

要求每个节点的答案，启发式合并没跑了，但是注意不要每次都遍历set

只需要计算贡献即可，因为插入的只有头部尾部和中间三种情况

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int> pll;
typedef pair<int,int> plll;
const int N=3e5+10;
const int inf=0x3f3f3f3f;
const int mod=998244353;
set<ll> s;
int h[N],ne[N],e[N],idx;
int son[N],sz[N],st[N];
ll ans[N];
int flag;
ll sum;
void add(int a,int b){
    e[idx]=b,ne[idx]=h[a],h[a]=idx++;
}
void dfs(int u,int fa){
    int i;
    sz[u]=1;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa)
            continue;
        dfs(j,u);
        sz[u]+=sz[j];
        if(sz[j]>sz[son[u]]){
            son[u]=j;
        }
    }
}
void cal(int u,int fa,int x){
    if(s.empty()){
        s.insert(u);
    }
    else{
        auto tmp=s.lower_bound(u);
        if(tmp==s.end()){
            tmp--;
            ll d=*(tmp);
            sum+=(u-d)*(u-d);
        }
        else if(tmp==s.begin()){
            ll d=*tmp;
            sum+=(u-d)*(u-d);
        }
        else{
            auto tmp1=--tmp;
            tmp++;
            ll d=(*(tmp)-*(tmp1));
            sum-=d*d;
            sum+=(u-*(tmp1))*(u-*tmp1);
            sum+=(u-*tmp)*(u-*tmp);
        }
        s.insert(u);
    }
    int i;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa||j==flag)
            continue;
        cal(j,u,x);
    }
}
void solve(int u,int fa,int keep){
    int i;
    for(i=h[u];i!=-1;i=ne[i]){
        int j=e[i];
        if(j==fa||j==son[u])
            continue;
        solve(j,u,0);
    }
    if(son[u]){
        solve(son[u],u,1);
        flag=son[u];
    }
    cal(u,fa,1);
    flag=0;
    ans[u]=sum;
    if(!keep){
        s.clear();
        sum=0;
    }
}
int main(){
    ios::sync_with_stdio(false);
    int i;
    int n;
    cin>>n;
    memset(h,-1,sizeof h);
    for(i=2;i<=n;i++){
        int p;
        cin>>p;
        add(i,p);
        add(p,i);
    }
    dfs(1,-1);
    solve(1,-1,0);
    for(i=1;i<=n;i++){
        cout<<ans[i]<<endl;
    }
}