dsu on tree——2020-camp-day2-E

/*
每个结点维护一个set存储该节点内的结点个数
然后找到最大的子树maxv，Smaxv拷贝给Su，再把u加入集合
然后遍历剩下子树，依次加进Su即可 
*/
#include<bits/stdc++.h>
using namespace std;
#define N 100005
#define ll long long 

vector<int>G[N];
set<ll>S[N];
set<ll>::iterator pre,nxt;

ll n,ans[N],p[N],size[N];

void dfs1(int u,int pre){
    size[u]=1;
    for(auto v:G[u])
        if(v==pre)continue;
        else {
            dfs1(v,u);
            size[u]+=size[v];
        }
}

void dfs(int u,int fa){
    if(G[u].size()==1 && u!=1){
        S[u].insert(u);return;
    }
    
    for(auto v:G[u]) 
        if(v!=fa)dfs(v,u);
    
    int Max=0,maxv;
    for(auto v:G[u]){
        if(v==fa)continue;
        if(Max<size[v]){
            Max=size[v];
            maxv=v;
        }
    }
    ans[u]=ans[maxv];
    S[u].swap(S[maxv]);
    S[maxv].clear();
    //把u加入S[u] 
    nxt=S[u].lower_bound(u);
    if(nxt==S[u].begin()){
        ans[u]+=1ll*(*nxt-u)*(*nxt-u);
    }
    else if(nxt==S[u].end()){
        pre=nxt;
        pre--;
        ans[u]+=1ll*(u-*pre)*(u-*pre);
    }
    else {
        pre=nxt;
        pre--;
        ans[u]+=1ll*(*nxt-u)*(*nxt-u);
        ans[u]+=1ll*(u-*pre)*(u-*pre);
        ans[u]-=1ll*(*nxt-*pre)*(*nxt-*pre);
    }
    S[u].insert(u);
    
    for(auto v:G[u]){
        if(v==maxv || v==fa)continue;
        for(auto x:S[v]){//把x并入S[u] 
            nxt=S[u].lower_bound(x);
            if(nxt==S[u].begin()){
                ans[u]+=1ll*(*nxt-x)*(*nxt-x);
            }
            else if(nxt==S[u].end()){
                pre=nxt;
                pre--;
                ans[u]+=1ll*(x-*pre)*(x-*pre);
            }
            else {
                pre=nxt;
                pre--;
                ans[u]+=1ll*(*nxt-x)*(*nxt-x);
                ans[u]+=1ll*(x-*pre)*(x-*pre);
                ans[u]-=1ll*(*nxt-*pre)*(*nxt-*pre);
            }
            S[u].insert(x);
        } 
        S[v].clear(); 
    }
} 

int main(){
    cin>>n;
    if(n==1){
        puts("0");
        return 0;
    }
    for(int i=2;i<=n;i++){
        scanf("%lld",&p[i]);
        G[i].push_back(p[i]);
        G[p[i]].push_back(i);
    }
    dfs1(1,0);
    dfs(1,0);
    
    for(int i=1;i<=n;i++)cout<<ans[i]<<'
';
}