【树形dp】Codeforces Round #405 (rated, Div. 1, based on VK Cup 2017 Round 1) B. Bear and Tree Jumps

时间:2023-03-09 09:31:58
【树形dp】Codeforces Round #405 (rated, Div. 1, based on VK Cup 2017 Round 1) B. Bear and Tree Jumps

我们要统计的答案是sigma([L/K]),L为路径的长度,中括号表示上取整。

[L/K]化简一下就是(L+f(L,K))/K,f(L,K)表示长度为L的路径要想达到K的整数倍,还要加上多少。

于是,我们现在只需要统计sigma((L+f(L,K))),最后除以K即可。

统计sigma(L)时,我们考虑计算每条边出现在了几条路径中,设u为edgei的子节点,那么这条边对答案的贡献就是siz(u)*(n-siz(u)),siz(u)为u的子树大小。

统计sigma(f(L,K))时,我们需要dp出f(u,r)表示从u的子树中出发,到u结束的路径中,长度mod K == r的有多少条。

具体统计方法跟

http://www.cnblogs.com/autsky-jadek/p/6580355.html

类似。

复杂度O(n*K^2)。

http://codeforces.com/blog/entry/51068

官方题解说得也很清楚。

#include<cstdio>
#include<iostream>
using namespace std;
typedef long long ll;
ll ans;
int v[400010],first[200010],next[400010],e;
ll f[200010][5];
void AddEdge(int U,int V){
v[++e]=V;
next[e]=first[U];
first[U]=e;
}
int n,m,fa[200010],siz[200010];
void dfs(int U){
f[U][0]=1;
siz[U]=1;
int tot=0;
for(int i=first[U];i;i=next[i]){
if(!fa[v[i]]){
fa[v[i]]=U;
dfs(v[i]);
siz[U]+=siz[v[i]];
ans+=(ll)siz[v[i]]*(ll)(n-siz[v[i]]);
for(int j=0;j<m;++j){
for(int k=0;k<m;++k){
ans+=(f[U][j]*f[v[i]][k])*(ll)((j+k+1)%m==0 ? 0 : m-(j+k+1)%m);
}
}
for(int j=0;j<m;++j){
f[U][(j+1)%m]+=f[v[i]][j];
}
}
}
}
int main(){
// freopen("c.in","r",stdin);
int x,y;
scanf("%d%d",&n,&m);
for(int i=1;i<n;++i){
scanf("%d%d",&x,&y);
AddEdge(x,y);
AddEdge(y,x);
}
fa[1]=-1;
dfs(1);
cout<<ans/(ll)m<<endl;
return 0;
}