Codeforces790B Bear and Tree Jumps -- 树形DP

2019-04-14 08:33发布

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

代码

#include #include #include #include #include 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;iif(g[x][i]!=y)Dfs(g[x][i],x); for(int j=0;jint p=j-1;if(p<0)p=k-1; for(int i=0;iint 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;iint v=g[x][i]; if(v==y)continue; for(int j=0;jfor(int p=0;p<2;p++)s[p][j]=f[p][v][j]; for(int j=0;jint 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;iif(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;iint 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;iscanf("%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<2<return 0; } Dfs(1,0);Dfs2(1,0); for(i=1;i<=n;i++) for(j=0;j1][i][j]; cout<2<return 0; }