题目
内容
假设子树的节点编号排序后的序列为
a
1
,
a
2
,
a
3
,
.
.
.
,
a
k
a_1,a_2,a_3,...,a_k
a1,a2,a3,...,ak ,这个节点的“结实程度”就是:
∑
i
=
1
k
−
1
(
a
i
+
1
−
a
i
)
2
\sum_{i=1}^{k-1}\left(a_{i+1}-a_i\right)^2
∑i=1k−1(ai+1−ai)2
求这棵树每个节点的结实程度.
分析
如果我们要求一个节点的结实程度,我们需要知道以这个节点为根的子树所有节点的编号.
其实也就是将当前节点的儿子节点的子树合并到一起,在加上当前节点本身.
还需要保证所有节点有序,因此我们用set来维护当前子树所有节点的编号.
但是这样复杂度是 O ( n 2 l o g n ) O(n^2logn) O(n2logn) ,所以我们还需要用到启发式合并的思想,也就是把size小的集合合并到size大的集合.具体来说:我们先找到当前节点的重儿子,把重儿子的集合作为当前节点的集合,再把其他儿子节点集合中的元素插进来.
复杂度 O ( n l o g 2 n ) O(nlog^2n) O(nlog2n)
#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;
}