2020 CCPC-Wannafly Winter Camp Day2 (Div.1&2) E. 阔力梯的树

启发式合并即可

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 2e5 + 10;
#define fi first
#define se second
#define pb push_back
vector<int>v[N];
int n,sz[N];
void dfs(int now){
  sz[now]=1;
  for(auto k:v[now]){
    dfs(k);
    sz[now]+=sz[k];
  }
}
LL ans[N],MAX=1e18,MIN=-1e18;
int num[N];
set<LL>fu[N];
void merge(int x,int y){
  for(auto k:fu[y]){
    if(k==MAX||k==MIN)
      continue;
    auto ne=--fu[x].upper_bound(k);
    auto f=fu[x].upper_bound(k);
    if(*f==MAX){
      if(*ne!=MIN){
        ans[x]+=(k-*ne)*(k-*ne);
      }
    }else{
      if(*ne==MIN){
        ans[x]+=(*f-k)*(*f-k);
      }else{
        ans[x]-=(*f-*ne)*(*f-*ne);
        ans[x]+=(*f-k)*(*f-k);
        ans[x]+=(k-*ne)*(k-*ne);
      }
    }
    fu[x].insert(k);
  }
}
LL ok[N];
void gao(int now){
  fu[now].insert(now);
  for(auto k:v[now]){
    gao(k);
    int x=num[now];
    int y=num[k];
    if(fu[x].size()>=fu[y].size()){
      merge(x,y);
      fu[y].clear();
    }else{
      merge(y,x);
      fu[x].clear();
      swap(num[k],num[now]);
    }
  }
  ok[now]=ans[num[now]];
}
int main() {
  scanf("%d",&n);
  for(int i=1;i<=n;i++){
    num[i]=i;
    fu[i].insert(MAX);
    fu[i].insert(MIN);
  }
  for(int i=2;i<=n;i++){
    int x;
    scanf("%d",&x);
    v[x].pb(i);
  }
  dfs(1);
  gao(1);
  for(int i=1;i<=n;i++){
    printf("%lld\n",ok[i]);
  }
  return 0;
}