【CF860E】Arkady and a Nobody-men
题意:给你一棵n个点的有根树。如果b是a的祖先,定义$r(a,b)$为b的子树中深度小于等于a的深度的点的个数(包括a)。定义$z(a)=\sum\limits r(a,b)$(b是a的祖先)。要你求出每个点的z值。
$n\le 5\times 10^5$
题解:一开始naive的思路:将所有点按深度排序,将深度相同的点统一处理,统计答案时相当于链加,链求和,用树剖+树状数组搞一搞,时间复杂度$O(n\log^2n)$。
后来看题解发现我这个想法简直菜爆了。我们先从树形DP的角度去想,先给出转移方程:
$ans(x)=ans(fa(x))+dep(x)+ans‘(x)$,ans‘(x)表示与a深度相同的点 对a的贡献。
现在问题变成了求ans‘,我们考虑在每个点对的lca处统计贡献。具体地,我们对于每个点x,维护若干个三元组(d,a,cnt)表示x的子树中有cnt个d级子孙,其中一个子孙为a。DP的过程就相当于在父亲节点处将所有儿子节点的三元组合并,在合并时顺便统计贡献。
具体地,合并方式如下:假如x有两个儿子,它们有三元组$(d,a,cnt_a)$和$(d,b,cnt_b)$,则:
1.ans‘(a)+=dep(x)\times cnt_b
2.ans‘(b)+=dep(x)\times cnt_a
3.得到新三元组(d,a,cnt_a+cnt_b)
但是后面的点 对b的贡献呢?我们发现后面的点 对a和b的贡献就是相同的了,所以我们建一个新图,在新图中从a到b连一条长度为$ans‘(b)-ans‘(a)$的边,最后在新图上DFS一下,最最后统计一下ans数组即可。
以上过程采用长链剖分优化,由于一开始的三元组个数为n,则每次合并都会减少一个三元组,所以时间复杂度O(n)。
#include <cstdio> #include <cstring> #include <iostream> #include <vector> using namespace std; const int maxn=500010; typedef long long ll; ll ans[maxn]; struct node { int v,x; node() {} node(int a,int b) {v=a,x=b;} }mem[maxn<<1],*f[maxn],*now=mem; int n,cnt,Cnt,rt; int to[maxn],nxt[maxn],head[maxn],dep[maxn],md[maxn],son[maxn],fa[maxn]; bool vis[maxn]; int To[maxn],Nxt[maxn],Head[maxn]; ll Val[maxn]; inline void add(int a,int b) { to[cnt]=b,nxt[cnt]=head[a],head[a]=cnt++; } inline void Add(int a,int b,ll c) { To[Cnt]=b,Val[Cnt]=c,Nxt[Cnt]=Head[a],Head[a]=Cnt++; } void dfs1(int x) { md[x]=0; for(int i=head[x];i!=-1;i=nxt[i]) { dep[to[i]]=dep[x]+1,dfs1(to[i]); if(md[to[i]]+1>md[x]) md[x]=md[to[i]]+1,son[x]=to[i]; } } void dfs2(int x) { if(f[x]==NULL) f[x]=now,now+=md[x]+2; if(son[x]) f[son[x]]=f[x]+1,dfs2(son[x]); f[x][0]=node(1,x); for(int i=head[x];i!=-1;i=nxt[i]) if(to[i]!=son[x]) { dfs2(to[i]); for(int j=0;j<=md[to[i]];j++) { node a=f[x][j+1],b=f[to[i]][j]; ans[b.x]+=1ll*dep[x]*a.v; ans[a.x]+=1ll*dep[x]*b.v; Add(a.x,b.x,ans[b.x]-ans[a.x]); f[x][j+1]=node(a.v+b.v,a.x); } } } void dfs3(int x) { for(int i=Head[x];i!=-1;i=Nxt[i]) ans[To[i]]=ans[x]+Val[i],dfs3(To[i]); } void dfs4(int x) { for(int i=head[x];i!=-1;i=nxt[i]) ans[to[i]]+=ans[x]+dep[x],dfs4(to[i]); } inline int rd() { int ret=0,f=1; char gc=getchar(); while(gc<‘0‘||gc>‘9‘) {if(gc==‘-‘) f=-f; gc=getchar();} while(gc>=‘0‘&&gc<=‘9‘) ret=ret*10+(gc^‘0‘),gc=getchar(); return ret*f; } int main() { n=rd(); int i; memset(head,-1,sizeof(head)),memset(Head,-1,sizeof(Head)); for(i=1;i<=n;i++) { fa[i]=rd(); if(!fa[i]) rt=i; else add(fa[i],i); } dep[rt]=1,dfs1(rt); dfs2(rt); for(i=0;i<=md[rt];i++) dfs3(f[rt][i].x); dfs4(rt); for(i=1;i<=n;i++) printf("%lld ",ans[i]); return 0; }