2021.6.1C树上dp

#include<iostream>
using namespace std;
const int N=2e5+10;
int n,k;
int ne[N<<1],hd[N<<1],to[N<<1],idx;
void add(int x,int y)
{
    ne[++idx]=hd[x];
    hd[x]=idx;
    to[idx]=y;
}
long long sum[N][7];
long long ans;
void solve(int top,int v)
{
    int l=0,r=0,u=sum[v][k-1];
    for(int i=0; i<=k-1; i++) l+=sum[top][i],r+=sum[v][i];
    for(int i=k-2; i>=0; i--) sum[v][i+1]=sum[v][i];
    sum[v][0]=u;
    sum[v][6]+=u;
    ans+=l*sum[v][6]+r*sum[top][6];
    for(int i=0; i<=k-1; i++)
        for(int j=0; j<=k-1; j++)
            ans+=1LL*((i+j)/k+((i+j)%k!=0))*sum[top][i]*sum[v][j];
    for(int i=0; i<=k-1; i++) sum[top][i]+=sum[v][i];
    sum[top][6]+=sum[v][6];

}
void dfs(int top,int fa)
{
    sum[top][0]=1;
    int i=hd[top];
    while(i)
    {
        int v=to[i];
        i=ne[i];
        if(v==fa) continue;
        dfs(v,top);
        solve(top,v);
    }
}
int main()
{
    cin>>n>>k;
    for(int i=1,x,y; i<=n-1; i++)
    {
        scanf("%d%d",&x,&y);
        add(x,y);
        add(y,x);
    }
    dfs(1,0);
    cout<<ans<<endl;
}