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
#include

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;ifor(int j=0;jint 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;iint main()
{
    //freopen("/home/yang/Desktop/r.txt","r",stdin);
    while(~scanf("%d%d",&n,&k)){
        init(n);
        for(int i=0;i1;++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;
}