题意:
给出一棵树,一个人可以在树上跳,每次最多跳(k(1 leq k leq 5))个点
定义(f(s,t))为从顶点(s)跳到顶点(t)最少需要跳多少次
求(sumlimits_{s<t}f(s,t))
分析:
注意到(k)很小,为了方便转移,定义:
- (sz(u,i))为(u)的子树中与(u)的距离模(k)等于(i)的孩子节点的个数
- (d(u,i))为从(u)跳到(sz(u,i))对应的节点的跳数之和
转移:
- (sz(u,(i+1)\, mod \, k) += sz(v, i)),其中(v)是(u)的儿子节点
- (d(u,(i+1)\, mod \, k) += d(v, i), 0<i<k)
- (d(u,1 \, mod \, k) += d(v,0) + sz(v, 0)),因为到(v)距离为(k)的整数被的那些节点转移到(u)时会多跳一步
统计答案:
在递归到点(u)时,我们统计那些(lca(s,t)=u)的路径的贡献
这样的路径有两类:
- (u)到其孩子节点的路径
- 来自(u)的两个不同子树的孩子节点形成的路径
第一类答案很好计算,就是(sum d(u,i))
在第二类答案时,假设枚举到(u)的某个子树(v),而且之前的子树的答案已经更新到到(d(u,i),sz(u,i))里面
枚举(i,j)分别为(v)之前的子树距离(u)模(k)为(i)的路径,和子树(v)中距离(u)模(k)为(j)的路径
两边各选一段路拼起来,可以通过公式快速求得答案,细节不再赘述
最终时间复杂度为(O(k^2n))
#include <cstdio>
#include <cstring>
#include <vector>
using namespace std;
typedef long long LL;
const int maxn = 200000 + 10;
vector<int> G[maxn];
int n, k, sz[maxn][5];
LL d[maxn][5], ans;
void dfs(int u, int fa) {
int tsz[5];
LL td[5];
for(int v : G[u]) if(v != fa) {
dfs(v, u);
memset(tsz, 0, sizeof(tsz));
for(int i = 0; i < k; i++) tsz[(i+1)%k] += sz[v][i];
memset(td, 0, sizeof(td));
for(int i = 1; i < k; i++) td[(i+1)%k] += d[v][i];
td[1%k] += d[v][0] + sz[v][0];
for(int i = 0; i < k; i++) {
for(int j = 0; j < k; j++) {
int delta = (i+j+k-1)/k - (i>0) - (j>0);
ans += d[u][i] * tsz[j] + td[j] * sz[u][i];
if(delta) ans += (LL) delta * sz[u][i] * tsz[j];
}
}
for(int i = 0; i < k; i++) { d[u][i] += td[i]; sz[u][i] += tsz[i]; }
}
for(int i = 0; i < k; i++) ans += d[u][i];
sz[u][0]++;
}
int main()
{
scanf("%d%d", &n, &k);
for(int i = 1; i < n; i++) {
int u, v; scanf("%d%d", &u, &v);
G[u].push_back(v);
G[v].push_back(u);
}
dfs(1, 0);
printf("%lld
", ans);
return 0;
}