HDU5593 ZYB's Tree

求 \(n\) 个点的树上对于每个点距离小于 \(k\) 的点的数量（边权均为 \(1\) ）。 \(n\leq 5\times 10^5，k\leq 10\) 。

设 \(f[u][i]\) 表示距离 \(u\) 点 \(i\) 距离以内并且在子树内的点，那么在 \(DFS\) 时可以用 \(O(nk)\) 的时间处理 \(f[u][i]=\sum f[v][i-1]\) ，其中 \(v\) 是 \(u\) 的孩子。

同理，设 \(g[u][i]\) 表示距离 \(u\) 点距离 \(i\) 以内并且在子树外（但包含自己）的点，这一项需要用父亲信息去更新儿子信息， \(g[u][i]=g[fa][i-1]\) ，同时可能会 \(u\rightarrow fa\rightarrow bro\) ，那么在 \(bro\) 的内部也可能会存在符合条件的点，所以 \(g[u][i]+=\sum f[bro][i-2]\) ，为了避免重复计算，可以在第一次 \(DFS\) 时在 \(fa\) 节点记录 \(sum[u][i]=\sum f[son][i]\) ，那么在调用时只用将 \(g[u][i]+=sum[fa][i-2]-f[u][i-2]\) 即可。

#include<bits/stdc++.h>
using namespace std;

int n,k,t;
struct edge{
    int to,nxt;
}e[500005];
int head[500005],ecnt;
void adde(int u,int v){
    e[++ecnt].to=v;
    e[ecnt].nxt=head[u];
    head[u]=ecnt;
}

int f[500005][15],g[500005][15],sum[500005][15];

void dfs1(int u,int fa){
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa)
            continue;
        dfs1(v,u);
        for(int j=1;j<=k;++j){
            f[u][j]+=f[v][j-1];
            sum[u][j]+=f[v][j];
        }
        sum[u][0]+=f[v][0];
    }
}

void dfs2(int u,int fa){
    if(fa){
        for(int i=1;i<=k;++i){
            g[u][i]+=g[fa][i-1];
        }
        for(int i=2;i<=k;++i){
            g[u][i]+=sum[fa][i-2]-f[u][i-2];
        }
    }
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to;
        if(v==fa)
            continue;
        dfs2(v,u);
    }
}



int main(){
    scanf("%d",&t);
    while(t--){
        scanf("%d %d",&n,&k);
        memset(head,0,sizeof head);
        memset(f,0,sizeof f);
        memset(g,0,sizeof g);
        memset(sum,0,sizeof sum);
        ecnt=0;
        long long A,B;
        scanf("%lld %lld",&A,&B);
        for(int i=2;i<=n;++i)
            adde((A*i+B)%(i-1)+1,i);
        for(int i=1;i<=n;++i){
            for(int j=0;j<=k;++j){
                f[i][j]=g[i][j]=1;
            }
        }
        dfs1(1,0);
        dfs2(1,0);
        int ans=0;
        for(int i=1;i<=n;++i)
            ans^=f[i][k]+g[i][k]-1;
        printf("%d\n",ans);
    }
    return 0;
}