[CF791D]Bear and Tree Jumps 树上DP
题面:
给定一颗大小为\(n\)的树,给定一个k,求\(\sum_{i=1}^n \sum_{j=1}^n \lceil \frac{dis(i,j)}{k} \rceil\)
范围&性质:\(1\le n \le 2\times 10^5,1\le k\le 5\)
分析:
朴素暴力做法:
对于每一个点dfs一遍,统计答案,复杂度\(O(n^2)\)
正解:(前排膜拜 Orz)
先考虑k=1怎么做,相当于求树上任意两点间路径和,直接枚举每条边,求出左右两部分内点的个数\(num1,num2\),相乘就好了,复杂度\(O(n)\)
那怎么向k>1扩展呢,问题在于k>1时路径长度除以k会有余数,那我们就考虑把余数补齐,答案就变成了 \(\frac{ans+\sum f(i,j)}{k}\),其中\(f(i,j)\)表示将\(dis(i,j)\)补成k的倍数需要的代价,这样就可以消掉上取整符号的影响了
那么问题就转化成了如何求\(\sum f(i,j)\),这时候就利用到k很小这一性质,我们记\(dp[i][j]\)表示u的子树内到根节点距离模\(k\)余\(j\)的点的个数,对于树上任意两点\(a,b\)而言,他们之间的距离\(dis(a,b)=dep[a]+dep[b]-2*dep[lca]\),\(f(a,b)=(dis(a,b)-k)mod \ k\),转移的时候我们就可以枚举\(u,v\)中dp状态第二维\(i,j\),这些点之间的距离\(dis=(i+j-dep[u]) mod\ k\)(因为\(u,v\)的\(lca\)是\(u\))单个点对需要的代价\(f=(dis-k)mod\ k\),由乘法原理得总的代价为\(f(i,j)\times dp[u][i]\times dp[v][j]\)
代码:
#include<bits/stdc++.h>
using namespace std;
namespace zzc
{
const int maxn = 200005;
int cnt=0,head[maxn];
long long dp[maxn][5],siz[maxn],ans=0;
int k,n;
struct edge
{
int to,nxt;
}e[maxn<<1];
void add(int u,int v)
{
e[++cnt].to=v;
e[cnt].nxt=head[u];
head[u]=cnt;
}
void dfs(int u,int fa,int dep)
{
dp[u][dep%k] = 1;siz[u]=1;
for (int i=head[u];i;i=e[i].nxt)
{
int v=e[i].to;
if (v==fa) continue;
dfs(v,u,dep+1);
for (int j = 0;j < k;j++)
{
for (int r = 0;r < k;r++)
{
int dis = (j + r - dep * 2) % k;
int rev = (k - dis) % k;
ans += rev*dp[u][j] * dp[v][r];
}
}
siz[u]+=siz[v];
for (int j = 0;j < k;j++) dp[u][j] += dp[v][j];
ans += (n - siz[v])*siz[v];
}
}
void work()
{
scanf("%d%d",&n,&k);
for (int i=1,x,y;i < n;i++)
{
scanf("%d%d",&x,&y);
add(x,y);add(y,x);
}
dfs(1, -1, 0);
printf("%lld\n",ans/k);
}
}
int main()
{
zzc::work();
return 0;
}