CodeForce VKcup C 树形dp

题意： 给出一棵树，一个人可以在树上跳，每次最多跳k(1≤k≤5)个点
定义f(s,t)为从顶点ss跳到顶点tt最少需要跳多少次
求∑(s<t)f(s,t)

链接： 点我

 

dp[i][j]表示以i点为根，子树节点到i距离的mod为i的节点个数

fs[i][j]表示以i点为根，子树节点到i距离的mod为i的距离/k的值

写出转移方程，求ans的过程就是两边树相乘，得到距离

#include<cstdio>
#include<iostream>
#include<algorithm>
#include<cstring>
#include<cmath>
#include<queue>
#include<map>
using namespace std;
#define MOD 1000000007
const int INF=0x3f3f3f3f;
const double eps=1e-5;
typedef long long ll;
#define cl(a) memset(a,0,sizeof(a))
#define ts printf("\n*****\n");
#define ptn printf("\n");
const int MAXN=200005;
vector<int> G[MAXN];
ll ans=0;
ll dp[MAXN][6];
ll fs[MAXN][6];
int n,k;
void dfs(int u,int pre){
    dp[u][0]=1;
    for(int i=0;i<G[u].size();i++){

        int v=G[u][i];
        if(v==pre)    continue;
        dfs(v,u);
        for(int j=0;j<=k;j++){
            for(int h=0;h<=k;h++){
                    ll temp=j+h+1;
                    if(temp%k) temp=temp/k+1;  //距离之和大于k，说明步数需要再+1
                    else temp=temp/k;
                ans+=temp*dp[u][j]*dp[v][h]+dp[u][j]*fs[v][h]+dp[v][h]*fs[u][j];
            }
        }
        for(int j=1;j<=k;j++){
            dp[u][j]+=dp[v][j-1],fs[u][j]+=fs[v][j-1];
        }
        dp[u][1]+=dp[v][k];
        fs[u][1]+=dp[v][k]+fs[v][k];
    }
}
int main()
{
    scanf("%d %d",&n,&k);
    int u,v;
    cl(dp);cl(fs);
    for(int i=0;i<n-1;i++){
        scanf("%d %d",&u,&v);
        G[u].push_back(v);
        G[v].push_back(u);
    }
    dfs(1,-1);
    printf("%I64d\n",ans);
}