CF771C Bear and Tree Jumps 題解

Problem

​ 有一顆\(n\)個結點的樹，一隻熊可以從當前節點可以跳到任何與當前節點距離不超過\(k\)的節點。定義\(f(u,v)\)為熊從\(u\)點到\(v\)點所需的最少跳躍次數，那麼，對於樹上的所有點對\((u,v)\)，\(f(u,v)\)的總和是多少。

Solution

​ 感覺這題做的人不多，而且很簡單，就寫篇題解

​ 看到\(k \leq 5\)，考慮設\(f_{u,i}\)表示在\(u\)的子樹內離\(u\)的距離模\(k\)為\(i\)的點到\(u\)的距離之和，再設\(g_{u,i}\)表示滿足上述條件的點數，易得轉移方程：

​ 對於\(i \ne 1\)

​ 對於\(i=1\)

​ 換根的時候直接加加減減即可

​ 還有，\(k=1\)時我寫的好像有點問題，會WA，所以我\(k=1\)時直接求了子樹大小

Code

#include<bits/stdc++.h>
#define LL long long
using namespace std;

LL Ans,sum,f[200005][6],g[200005][6];

int n,m,cnt;

int size[200005];

int head[200005],to[400005],Next[400005];

inline int read(){
    int x=0,f=1;char ch=getchar();
    while(ch<'0'||ch>'9'){
       if(ch=='-')f=-1;
       ch=getchar();
    }
    while(ch>='0'&&ch<='9'){
       x=(x<<1)+(x<<3)+ch-'0';
       ch=getchar();
    }
    return x*f;
}

inline void add(int u,int v){
    to[++cnt]=v;Next[cnt]=head[u];head[u]=cnt;
}

void Dfs1(int u,int fa){
	size[u]=1;
	for(register int i=head[u];i;i=Next[i]){
		int v=to[i];
		if(v==fa)
			continue;
		Dfs1(v,u);
		size[u]+=size[v];
	}	
	if(u!=1)
		sum+=size[u];
	return;
}

void Dfs2(int u,int fa){
	Ans+=sum;
	for(register int i=head[u];i;i=Next[i]){
		int v=to[i];
		if(v==fa)
			continue;
		size[u]-=size[v];sum-=size[v];
		size[v]+=size[u];sum+=size[u];
		Dfs2(v,u);
		size[v]-=size[u];sum-=size[u];
		size[u]+=size[v];sum+=size[v];
	}
	return;
}

void DP1(int u,int fa){
    for(register int i=head[u];i;i=Next[i]){
        int v=to[i];
        if(v==fa) 
        	continue;
        DP1(v,u);
        for(register int k=1;k< m;++k){
        	f[u][(k+1)%m]+=f[v][k];
        	g[u][(k+1)%m]+=g[v][k];
		}		
		g[u][1]+=1+g[v][0];
		f[u][1]+=1+g[v][0]+f[v][0];
	}
	return; 
}

void DP2(int u,int fa){
	for(register int i=0;i< m;++i)
		Ans+=f[u][i];
	for(register int i=head[u];i;i=Next[i]){
		int v=to[i];
		if(v==fa)
			continue;
		for(register int k=1;k< m;++k){
			f[u][(k+1)%m]-=f[v][k];
			g[u][(k+1)%m]-=g[v][k];
		}
		g[u][1]-=1+g[v][0];
		f[u][1]-=1+g[v][0]+f[v][0];
        for(register int k=1;k< m;++k){
        	f[v][(k+1)%m]+=f[u][k];
        	g[v][(k+1)%m]+=g[u][k];
		}		
		g[v][1]+=1+g[u][0];
		f[v][1]+=1+g[u][0]+f[u][0];		
		DP2(v,u);		
		g[v][1]-=1+g[u][0];
		f[v][1]-=1+g[u][0]+f[u][0];
        for(register int k=1;k< m;++k){
        	f[v][(k+1)%m]-=f[u][k];
        	g[v][(k+1)%m]-=g[u][k];
		}		
		g[u][1]+=1+g[v][0];
		f[u][1]+=1+g[v][0]+f[v][0];		
		for(register int k=1;k< m;++k){
			f[u][(k+1)%m]+=f[v][k];
			g[u][(k+1)%m]+=g[v][k];
		}		
	}	
	return;
}
        
int main(){
    
    n=read();m=read();

    for(register int i=2;i<=n;++i){
        int u,v;
        u=read();v=read();
        add(u,v);add(v,u);
    }

    if(m==1){
    	Dfs1(1,0);
		Dfs2(1,0);
	}
	else{
    	DP1(1,0);
		DP2(1,0);
	}

	printf("%lld\n",Ans>>1LL);

    return 0;
}