CF771C Bear and Tree Jumps 题解

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\)

\[f_{u,i}=\sum\limits_{v \in son}f_{v,(i-1+k)\%k} \\ g_{u,i}=\sum\limits_{v \in son}g_{v,(i-1+k)\%k} \]

​ 对于\(i=1\)

\[g_{u,1}=\sum\limits_{v \in son}g_{v,0}+1 \\ f_{u,1}=\sum\limits_{v \in son}f_{v,0}+g_{v,0}+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;
}