VK Cup 2017 - Round 1 C. Bear and Tree Jumps 换根dp

https://www.luogu.com.cn/blog/Hakurei-Reimu/solution-cf771c 参考题解

#include<bits/stdc++.h>

using namespace std;

#define X first
#define Y second
#define eps  1e-2
#define gcd __gcd
#define pb push_back
#define PI acos(-1.0)
#define lowbit(x) (x)&(-x)
#define bug printf("!!!!!\n");
#define mem(x,y) memset(x,y,sizeof(x))

#define int long long 
typedef long long LL;
typedef long double LD;
typedef pair<int,int> pii;
typedef unsigned long long uLL;

const int maxn = 2e5+2;
const int INF  = 1<<30;
const int mod  = 1e9+7;

int n,k,f[maxn][6],sz[maxn];
std::vector<int> v[maxn];
void dfs1(int x,int fa){
    sz[x]=1;
    for(int i=0;i<v[x].size();i++){
        int to=v[x][i];if(to==fa) continue;
        dfs1(to,x);
        sz[x]+=sz[to];
        f[x][0]+=f[to][k-1]+sz[to];
        for(int j=1;j<k;j++) f[x][j]+=f[to][j-1];
    }
}
void dfs2(int x,int fa){
    LL t[5]={0};
    if(fa){
        t[0]=f[fa][0]-f[x][k-1]-sz[x];
        for(int i=1;i<k;i++) t[i]=f[fa][i]-f[x][i-1];
    }
    f[x][0]+=t[k-1]+n-sz[x];   
    for(int i=1;i<k;i++) f[x][i]+=t[i-1];
    for(int i=0;i<v[x].size();i++){
        int to=v[x][i];if(to==fa) continue;
        dfs2(to,x);
    }
}
void solve(){
    scanf("%lld%lld",&n,&k);
    for(int i=1;i<n;i++){
        int x,y;scanf("%lld%lld",&x,&y);
        v[x].push_back(y);
        v[y].push_back(x);
    }
    dfs1(1,0);
    dfs2(1,0);
    LL ans=0;for(int i=1;i<=n;i++) ans+=f[i][0];
    cout<<ans/2<<endl;
    return;
}

int32_t main()
{
//    freopen("in.txt","r",stdin);
//    freopen("out.txt","w",stdout);
//    ios::sync_with_stdio(false);
    int t = 1;
    //scanf("%d",&t);
    while(t--){
    //    printf("Case %d: ",cas++);
        solve();
    }
    return 0;
}