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_back

typedef 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);
  }
}