Codeforces790B. Bear and Tree Jumps 【树型DP】

传送门

---

$ans=\sum{ceiling(\frac{dis_{i,j}}{k})}$
$令f(x,k)=k-x\%k$
$则ans=\sum{ \frac{ dis_{i,j}+f(dis_{i,j},k) }{k} } = \frac1k(\sum{ dis_{i,j} } + \sum{ f(dis_{i,j},k) })$

$于是问题转化为求\sum{ dis_{i,j} }和 \sum{ f(dis_{i,j},k) }$

---

$令sz[i]=以i为根的子树包含的节点个数,far[i]为i的父亲节点$
$则边{far[i],i}对\sum{dis_{i,j}}的贡献=sz[i]*(n-sz[i])$
$\sum{dis_{i,j}}=\sum{sz[i](n-sz[i])}$

---

$令d[u][i]=以u为根的子树,深度\%k=i的节点个数$
$若v为u的儿子,u深度为dept$
$d[u][i]=以i为根的子树,深度\%k=i的节点个数(不含有v为根的子树包含的点)$
$距离为dis=(i+j-2*dept)\%k的点对个数为d[u][i]*d[v][j]$
$k-dis则为f(dis,k),对\sum{f(dis_{a,b},k)}的贡献为(k-dis)*d[u][i]*d[v][j]$

#include<stdio.h>
#include<bits/stdc++.h>

using namespace std;

#define ll long long
#define MEM(a,x) memset(a,x,sizeof(a))

const int N = 2e5+5;

vector<int>G[N];
int sz[N],d[N][10];
ll ans;
int n,k;

inline int mod(int x){
    return (x%k+k)%k;
}

void init(int n){
    for(int i=0;i<=n;++i){
        G[i].clear();
        sz[i]=0;
        MEM(d[i],0);
    }
    ans=0;
}

void dfs(int u,int far,int dept){
    sz[u]=d[u][dept%k]=1;
    for(int v:G[u]){
        if(v==far){
            continue;
        }
        dfs(v,u,dept+1);
        for(int i=0;i<k;++i){
            for(int j=0;j<k;++j){
                int dis=mod(i+j-2*dept);
                ans+=(ll)mod(k-dis)*d[u][i]*d[v][j];
            }
        }
        sz[u]+=sz[v];
        for(int i=0;i<k;++i){
            d[u][i]+=d[v][i];
        }
    }
    ans+=(ll)sz[u]*(n-sz[u]);
}

int main()
{
    //freopen("/home/yang/Desktop/r.txt","r",stdin);
    while(~scanf("%d%d",&n,&k)){
        init(n);
        for(int i=0;i<n-1;++i){
            int u,v;
            scanf("%d%d",&u,&v);
            G[u].push_back(v);
            G[v].push_back(u);
        }
        dfs(1,-1,0);
        printf("%lld\n",ans/k);
    }
    return 0;
}