[YZOJ]P2642-主席树-统计子树信息

链接
利用DFS序的时间上的性质，我们可以将访问后得到的深度减去访问前得到的深度来统计答案。前缀和即可处理。

#include<cstdio>
#include<algorithm>

#define R register
#define max_n 1000010
using namespace std;
int n,dis[max_n];
int lson[max_n],rson[max_n];
int dep[max_n],cnt[max_n];
int lans[max_n],rans[max_n];

int read()
{
    R int xx;R char ch;
    while(ch=getchar(),ch<'0'||ch>'9');xx=ch-'0';
    while(ch=getchar(),ch>='0'&&ch<='9')xx=(xx<<1)+(xx<<3)+ch-'0';
    return xx;
}

void dfs(R int u)
{
    cnt[dep[u]]++;
    R int k=dep[u]+dis[u];
    if(lson[u])
    {
        lans[u]=-cnt[k];
        dep[lson[u]]=dep[u]+1,dfs(lson[u]);
        lans[u]+=cnt[k];
    }
    if(rson[u])
    {
        rans[u]=-cnt[k];
        dep[rson[u]]=dep[u]+1,dfs(rson[u]);
        rans[u]+=cnt[k];
    }
}

void print()
{
    register const int P = 1004535809, Q = 998244353;
    register int A = 2341, B = 4123;
    for(register int i = 1; i <= n; ++i)
    {
        A = ((long long) A * n + lans[i]) % P; B = ((long long) B * n + lans[i]) % Q;
        A = ((long long) A * n + rans[i]) % P; B = ((long long) B * n + rans[i]) % Q;
    }
    printf("%d %d\n", A, B);
}


int main()
{
    R int i,j;
    R int u,v,et=0;
    n=read();
    for(i=1;i<=n;++i)
        dis[i]=read();
    for(i=1;i<=n;++i)
        lson[i]=read(),rson[i]=read();
    dep[1]=1,dfs(1);
    print();
    return 0;
}