codeforces 405 D. Bear and Tree Jumps 树形dp

题目链接

http://codeforces.com/contest/791/problem/D

题意

给一个树，节点数为n(2<=n<=2e5)，给一个数k(1<=k<=5)，代表一个人一步可走的最大距离，定义函数f(s,t)为树上s，t两点之间的最短步数，求∑ni=1∑nj=i+1f(i,j)。

思路

官方题解在这里：http://codeforces.com/blog/entry/51068（当然如果我看的懂的话，现在就不写博客了QAQ）
首先对于k=1，或者说直接求树上两点间距离之和，就是∑ni=1size[i]∗(n−size[i])，对于k不等于1的情况，麻烦的是如果L%k!=0，那么对于L，还要加上k−(L%k) 才能被k整除，k==1 时的ans很容易求，我们只要使答案加上所有的偏置量，再除以k就好了。
如何求所有的偏置量呢？用树形dp，对于每个节点u，我们分别计算以u为根节点的子树中，一定经过u的路径的总偏置量之和，不断递归下去，就可以求得所有偏置量。
对于偏置量的具体求法，我们记录每个节点到root的距离，树中任意一条路径的长度L=depth[u]+depth[v]−2∗depth[fa] ，其中fa是u，v的公共祖先。鉴于k的范围很小（k的范围必须很小，否则就不是这个档次的题了23333333），进行k2 大小的枚举，计算出当前偏置量的大小，并依据计数来统计有几个这样的路径。
计数统计的部分很简单，就不多说了，详见代码。

代码

#include <cstdio>#include <cstring>#include <iostream>#include <algorithm>#include <vector>using namespace std;#define MS(x, y) memset(x, y, sizeof(x))#define PB push_backtypedef long long LL;const int MAXN = 2e5 +5;vector<int> vec[MAXN];int n, k;LL ans;LL cnt[MAXN][6], siz[MAXN];void dfs(int u, int fa, int dep) {  cnt[u][dep % k] = siz[u] = 1;  int v, need;  for (int i = 0; i < vec[u].size(); ++i) {    v = vec[u][i];    if (v == fa) continue;    dfs(v, u, dep + 1);    for (int j = 0; j < k; ++j) for (int l = 0; l < k; ++l) {      need = ((k - (j + l - 2 * dep)) % k + k) % k;      ans += 1ll * need * cnt[u][j] * cnt[v][l];    }    for (int j = 0; j < k; ++j) cnt[u][j] += cnt[v][j];    siz[u] += siz[v];  }  ans += 1ll * siz[u] * (n - siz[u]);}int main() {  while (~scanf("%d%d", &n, &k)) {    for (int i = 1; i <= n; ++i) vec[i].clear();    MS(cnt, 0);    MS(siz, 0);    for (int i = 1, u, v; i < n; ++i) {      scanf("%d%d", &u, &v);      vec[u].PB(v);      vec[v].PB(u);    }    ans = 0;    dfs(1, 0, 0);    printf("%I64d\n", ans / k);  }}