老套路。。。
令
gi,j
表示
i
的子树中到
代码
#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
using namespace std;
#define N 200010
#define ll long long
vector<int>g[N];
int x,y;
ll Ans,f[2][N][5],s[2][5];
int i,j,k,n,m,p;
inline void Dfs(int x,int y){
f[0][x][0]=1;
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y)Dfs(g[x][i],x);
for(int j=0;j<k;j++){
int p=j-1;if(p<0)p=k-1;
for(int i=0;i<g[x].size();i++){
int v=g[x][i];
if(v==y)continue;
f[0][x][j]+=f[0][v][p];f[1][x][j]+=f[1][v][p]+(!p?f[0][v][0]:0);
}
}
}
inline void Dfs2(int x,int y){
for(int i=0;i<g[x].size();i++){
int v=g[x][i];
if(v==y)continue;
for(int j=0;j<k;j++)
for(int p=0;p<2;p++)s[p][j]=f[p][v][j];
for(int j=0;j<k;j++){
int p=j-1;if(p<0)p=k-1;
f[1][v][j]+=f[1][x][p]-(s[1][(!p?k-1:p-1)]+(p==1?s[0][0]:0))+(!p?f[0][x][0]-s[0][k-1]:0);
f[0][v][j]+=f[0][x][p]-s[0][(!p?k-1:p-1)];
}
Dfs2(v,x);
}
}
inline void Dfs3(int x,int y){
f[0][x][0]=1;
for(int i=0;i<g[x].size();i++)
if(g[x][i]!=y){
Dfs3(g[x][i],x);
f[0][x][0]+=f[0][g[x][i]][0];
f[1][x][0]+=f[1][g[x][i]][0]+f[0][g[x][i]][0];
}
}
inline void Dfs4(int x,int y){
for(int i=0;i<g[x].size();i++){
int v=g[x][i];
if(v==y)continue;
f[1][v][0]+=f[1][x][0]-(f[1][v][0]+f[0][v][0])+f[0][x][0]-f[0][v][0];
f[0][v][0]+=f[0][x][0]-f[0][v][0];
Dfs4(v,x);
}
}
int main(){
scanf("%d%d",&n,&k);
for(i=1;i<n;i++)scanf("%d%d",&x,&y),g[x].push_back(y),g[y].push_back(x);
if(k==1){
Dfs3(1,0);Dfs4(1,0);
for(i=1;i<=n;i++)Ans+=f[1][i][0];
cout<<Ans/2<<endl;
return 0;
}
Dfs(1,0);Dfs2(1,0);
for(i=1;i<=n;i++)
for(j=0;j<k;j++)
Ans+=f[1][i][j];
cout<<Ans/2<<endl;
return 0;
}