【树型DP】codeforces 791d

对于树型DP还根本一窍不通的我来说，竟然在div2的D里遇到了。。一脸尴尬啊

题目：http://codeforces.com/contest/791/problem/D
题意：一颗树，现在你能一次最多跳k格，问所有节点相互之间访问所用的步骤数是多少。

节点数为200000，不可能暴力去遍历，这时候要用到DP，但是怎么DP呢，首先要搞清楚题目到底需要你干啥。
对于所有节点每次只走一步，每两个节点之间的步骤数可以写成ak+b，表示跳a次之后，还需要走b步，那么你只需要找需要补多少步b的个数即可
所以可以构建状态方程dp[root][k]，表示当前节点你可以走出来的1到k的步骤数为那么多。
那么对于每一个父节点则有dp[root][j] = Σdp[v[root][i]][(j - 1 + k) % k]
为了防止重复统计，每次先计算当前步骤数的组合数为T[(i + j + 1) % k] += dp[root][i] * dp[g[root][p]][j];

最后只需要统计所有T[i] * (k - i) 即可

代码：

/*
@resources: codeforces 791D
@date: 2017-3-19
@author: QuanQqqqq
@algorithm: dfs 树型dp
*/
#include <bits/stdc++.h>

#define MAXN 200005
#define ll long long

using namespace std;

vector<ll> g[MAXN];
ll dp[MAXN][6];
ll T[6];
ll n,k,avg;

int dfs(int root,int pre){
    dp[root][0] = 1;
    ll r = 1;
    int l = g[root].size();
    for(int p = 0;p < l;p++){
        if(g[root][p] != pre){
            ll tmp = dfs(g[root][p],root);
            r += tmp;
            avg += 1LL * tmp * (n - tmp);
            for(int i = 0;i < k;i++){
                for(int j = 0;j < k;j++){
                    T[(i + j + 1) % k] += dp[root][i] * dp[g[root][p]][j];
                }
            }
            for(int i = 0;i < k;i++){
                dp[root][i] += dp[g[root][p]][(i + k - 1) % k];
            }
        }
    }
    return r;
}

int main(){
    int u,v;
    scanf("%lld %lld",&n,&k);
    for(int i = 0;i < n - 1;i++){
        scanf("%d %d",&u,&v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    avg = 0;
    dfs(1,0);
    for(int i = 1;i < k;i++){
        avg += 1LL * (k - i) * T[i];
    }
    printf("%lld\n",avg / k);
}