Codeforces790B Bear and Tree Jumps -- 树形DP

老套路。。。
令 gi,j 表示 i 的子树中到 i 距离模 k 为 j 的结点数， fi,j 表示 i 的子树中到 i 距离模 k 为 j 的距离和。然后就可以转移了。

代码

#include<cstdio>#include<vector>#include<cstring>#include<iostream>#include<algorithm>using namespace std;#define N 200010#define ll long longvector<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;}