Codeforces790B. Bear and Tree Jumps 【树型DP】

传送门

---

ans=∑ceiling(disi,jk)
令f(x,k)=k−x%k
则ans=∑disi,j+f(disi,j,k)k=1k(∑disi,j+∑f(disi,j,k))

于是问题转化为求∑disi,j和∑f(disi,j,k)

---

令sz[i]=以i为根的子树包含的节点个数,far[i]为i的父亲节点
则边far[i],i对∑disi,j的贡献=sz[i]∗(n−sz[i])
∑disi,j=∑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),对∑f(disa,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;}