HDU 5593 树形DP

HDU 5593
题目链接：
http://acm.hdu.edu.cn/showproblem.php?pid=5593
题意：
给含有n个点(n <= 5 * 1e5)的一棵树，1为根节点。
Ans[i]表示树中到i的距离小于等于k的点的个数。求所有ans[i]的异或和。
思路：
现场的时候并没有看到这道题。
很容易想到是O(n)的做法，所以只能遍历树一遍。然后发现k特别小，所以dp[i][j]来树形dp是可以承受的。
那么具体的话，dp[i][j]表示已i为根节点的子树，距离i距离不超过j的节点个数。注意是不超过，具体处理办法可以先把距离为j的先求出来，然后求前缀和。
然后对于每个节点的合法值，dp[u][k] + (dp[fa[u][k-1]] - dp[u][k-2]) + … + dp[v][0]。解释的话就是它的父亲结点的距离不超过k-1且不是从u的距离不是k-2的点中得到的点。依次类推。由于每个点最多找k次，所以暴力查找就可以。
源码：

#include <cstdio>
#include <cstring>
#include <cstdlib>
#include <cmath>
#include <algorithm>
#include <iostream>
#include <string>
#include <vector>
using namespace std;
#define LL long long
const int MAXN = 500000 + 5;
int fa[MAXN], f[MAXN][11];
vector<int>lin[MAXN];
int ans[MAXN];
int n, k;
void dfs1(int u)
{
    for(int i = 0 ; i < (int)lin[u].size() ; i++){
        int v = lin[u][i];
        dfs1(v);
        for(int j = 1 ; j <= k ; j++)   f[u][j] += f[v][j - 1];
    }
    f[u][0] = 1;
//    for(int i = 1 ; i <= k ; i++)   f[u][i] += f[u][i - 1];
}
void dfs2(int u)
{
//    ans[u] = 0;
    for(int i = 0 ; i < (int)lin[u].size() ; i++){
        int v = lin[u][i];
        dfs2(v);
    }
    ans[u] = f[u][k];
//    for(int i = 0 ; i <= k ; i++)   ans[u] += f[u][i];
    int now = u;
    int cnt = 1;
    while(fa[now] != 0 && cnt <= k){
        ans[u] += f[fa[now]][k - cnt];
        if(k - cnt - 1 >= 0)    ans[u] -= f[now][k - cnt - 1];
        cnt++;
        now = fa[now];
    }
}
int main()
{
    int T;
    scanf("%d", &T);
    while(T--){
        LL a, b;
        scanf("%d%d%I64d%I64d", &n, &k, &a, &b);
        for(int i = 1 ; i <= n ; i++)   lin[i].clear();
        fa[0] = 0;
        fa[1] = 0;
        for(int i = 2 ; i <= n ; i++){
            fa[i] = (a * i + b) % (i - 1) + 1;
            lin[fa[i]].push_back(i);
        }
        for(int i = 1 ; i <= n ; i++){
            for(int j = 0 ; j <= k ; j++){
                f[i][j] = 0;
            }
        }
        dfs1(1);
        for(int i = 1 ; i <= n ; i++){
            for(int j = 1 ; j <= k ; j++){
                f[i][j] += f[i][j - 1];
            }
        }
        dfs2(1);
        int res = 0;
        for(int i = 1 ; i <= n ; i++)   res = res ^ ans[i];
//        for(int i = 1 ; i <= n ; i++)   printf("ans[%d] = %d\n", i, ans[i]);
        printf("%d\n", res);
    }
    return 0;
}