Wannafly Camp 2020 Day 2E 阔力梯的树 - set,启发式合并

搞一波启发式合并即可

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define iter set<long long>::iterator

const int N = 100005;

struct myset {
    set <int> s;
    int ans = 0;
    void insert(int x) {
        iter p = s.insert(x).first;
        iter rb = s.end();
        --rb;
        if(p!=s.begin() && p!=rb) {
            iter pre=p, suf=p;
            --pre; ++suf;
            ans += x*x*2;
            ans += (*pre) * (*suf) * 2;
            ans -= x * ((*pre) + (*suf)) * 2;
        }
        if(p!=s.begin() && p==rb) {
            iter pre=p;
            --pre;
            ans += x*x + (*pre)*(*pre) - 2*x*(*pre);
        }
        if(p==s.begin() && p!=rb) {
            iter suf=p;
            ++suf;
            ans += x*x + (*suf)*(*suf) - 2*x*(*suf);
        }
    }
    void merge(myset *ms) {
        for(iter it=ms->s.begin();it!=ms->s.end();it++) {
            insert(*it);
        }
    }
    void print() {
        for(iter it=s.begin();it!=s.end();it++)
            cout<<(*it)<<" ";
    }
} buf[N];

myset* merge(myset *s1,myset *s2) {
    if(s1->s.size() < s2->s.size()) {
        s2->merge(s1);
        return s2;
    }
    else {
        s1->merge(s2);
        return s1;
    }
}

int n,p[N],vis[N],ans[N];
myset *s[N];
vector <int> g[N];

void print() {
    for(int i=1;i<=n;i++) {
        cout<<"Node "<<i<<": ";
        s[i]->print();
        cout<<endl;
    }
}

void dfs(int p) {
    vis[p]=1;
    for(int i=0;i<g[p].size();i++) {
        int q=g[p][i];
        if(vis[q]) continue;
        dfs(q);
        s[p]=merge(s[p],s[q]);
    }
    ans[p]=s[p]->ans;
}

signed main() {
    ios::sync_with_stdio(false);
    cin>>n;
    for(int i=1;i<=n;i++) {
        s[i]=&buf[i];
        s[i]->insert(i);
    }
    for(int i=2;i<=n;i++) {
        cin>>p[i];
        g[p[i]].push_back(i);
        g[i].push_back(p[i]);
    }
    dfs(1);
    for(int i=1;i<=n;i++) {
        cout<<ans[i]<<endl;
    }
}