题目链接
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);
}
}