阔力梯的树-启发式合并

题目

题目链接

内容

假设子树的节点编号排序后的序列为$a_1,a_2,a_3,...,a_k$ ，这个节点的“结实程度”就是：
${\sum }_{i=1}^{k-1}{\left(a_{i+1}-a_i\right)}^2$

求这棵树每个节点的结实程度.

分析

如果我们要求一个节点的结实程度,我们需要知道以这个节点为根的子树所有节点的编号.

其实也就是将当前节点的儿子节点的子树合并到一起,在加上当前节点本身.

还需要保证所有节点有序,因此我们用set来维护当前子树所有节点的编号.

但是这样复杂度是$O(n^{2}logn)$ ,所以我们还需要用到启发式合并的思想,也就是把size小的集合合并到size大的集合.具体来说:我们先找到当前节点的重儿子,把重儿子的集合作为当前节点的集合,再把其他儿子节点集合中的元素插进来.

复杂度$O(nlo{g}^{2}n)$

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
typedef pair<int,int> pii;
typedef vector<int> vi;
#define debug(x) cerr<<#x<<' '<<x<<'\n'
#define mp make_pair
#define pb push_back
#define fi first
#define se second
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define per(i,a,b) for(int i=(a);i>=(b);i--)
const int maxn=1e5+10;
const int mod=1e9+7;
const int inf=0x3f3f3f3f;
vi g[maxn];
set<int> s[maxn];
int cnt=0,f[maxn];
int fa[maxn],siz[maxn];
ll ans[maxn];
ll pow2(ll a,ll b) {
    return (a-b)*(a-b);
}
void dfs(int now) {
    siz[now]=1;
    int id=0,Max=0;
    for(int x:g[now]) {
        if(x==fa[now]) continue;
        dfs(x);
        siz[now]+=siz[x];
        if(siz[x]>Max) id=x,Max=siz[x];
    }
    if(siz[now]==1) {
        ans[now]=0;
        f[now]=cnt++;
        s[f[now]].insert(now);
        return ;
    }
    f[now]=f[id];
    ans[now]=ans[id];
    for(int x:g[now]) {
        if(x==fa[now]||x==id) continue;
        for(int xx:s[f[x]]) {
            auto it=s[f[now]].lower_bound(xx);
            auto it2=it;
            it2--;
            if(it!=s[f[now]].end()) ans[now]+=pow2(*it,xx);
            if(it!=s[f[now]].begin()) ans[now]+=pow2(*(it2),xx);
            if(it!=s[f[now]].end()&&it!=s[f[now]].begin()) ans[now]-=pow2(*it,*(it2));
            s[f[now]].insert(xx);
        }
    }
    auto it=s[f[now]].lower_bound(now);
    auto it2=it;
    it2--;
    if(it!=s[f[now]].end()) ans[now]+=pow2(*it,now);
    if(it!=s[f[now]].begin()) ans[now]+=pow2(*it2,now);
    if(it!=s[f[now]].end()&&it!=s[f[now]].begin()) ans[now]-=pow2(*it,*it2);
    s[f[now]].insert(now);
}
int main()
{
    ios::sync_with_stdio(false);cin.tie(0);
    int n;
    cin>>n;
    rep(i,2,n) {
        cin>>fa[i];
        g[i].pb(fa[i]);
        g[fa[i]].pb(i);
    }
    dfs(1);
    rep(i,1,n) cout<<ans[i]<<'\n';
    return 0;
}