CF771C Bear and Tree Jumps

给定一棵 n 个节点的树与正整数 k，求

\sum_{i=1}^n\sum_{j=1}^n\left\lceil\dfrac{\text{dist}(i,j)}{k}\right\rceil

的值。1\le n\le 2\times 10^5,1\le k\le 5。

尽管这题可以有 O(nk^2) 的 dp 做法，但是仅限于边权为 1 的情况。

而点分治可以处理边权任意的情形！考虑当前根为 \text{rt} 的时候怎么算。

对每个 u 处理出来 \text{rt}\to u 路径上的边权和，记作 d_u，那么要求的就是：

其中 b_u 表示 u 是从哪个子树过来的。我们要能快速算出这个东西就行了。

注意到 k 很小，考虑将 d_u 按照 \bmod k 的余数分类。

由于 d_u+d_v=k\lfloor d_u/k\rfloor+(d_u\bmod k)+k\lfloor d_v/k\rfloor +(d_v\bmod k)，故

前一项可以直接算，考虑后一项怎么算。

设 f_i 为当前子树内满足 d_u\bmod k=i 的 u 的个数，g_i 为前面子树内 \bmod k=i 的个数，那么后一项的和就是

维护 f,g 以及 f 的前缀和即可。O(nk\log n)。

感觉好像可以做到 O((n+k)\log n) 之类的？

#include<bits/stdc++.h>

#define int long long

using namespace std;

inline int read(){
    int x=0,f=1;char c=getchar();
    for(;(c<'0'||c>'9');c=getchar()){if(c=='-')f=-1;}
    for(;(c>='0'&&c<='9');c=getchar())x=x*10+(c&15);
    return x*f;
}

const int MN=2e5+5;
const int MK=10;

int f[MK],g[MK],fs[MK],n,k;
vector<int>G[MN];
int sz[MN],rt,all,mx[MN],ans=0;
bool vis[MN];

void calcsiz(int u,int fa){
    mx[u]=0,sz[u]=1;
    for(int v:G[u]){
        if(vis[v]||v==fa)continue;
        calcsiz(v,u),sz[u]+=sz[v],mx[u]=max(mx[u],sz[v]);
    }
    mx[u]=max(mx[u],all-sz[u]);
    if(mx[u]<mx[rt]||rt==0)rt=u;
}

int now=0;

void calcdis(int u,int fa,int de){
    f[de%k]++,now+=(int)(de/k);
    for(int v:G[u]){
        if(v==fa||vis[v])continue;
        calcdis(v,u,de+1);
    }
}

void work(int u){
    memset(g,0,sizeof(g)),memset(f,0,sizeof(f)),memset(fs,0,sizeof(fs));
    vis[u]=1,g[0]=1;int cnt=1,sum=0;
    for(int v:G[u]){
        if(vis[v])continue;
        memset(f,0,sizeof(f)),now=0,calcdis(v,u,1);
        ans+=(now*cnt+sum*sz[v]),sum+=now,cnt+=sz[v];
        fs[0]=f[0];for(int i=1;i<k;i++)fs[i]=fs[i-1]+f[i];
        for(int i=0;i<k;i++){
            if(i==0)ans+=g[i]*(fs[k-1]-fs[0]);
            else ans+=g[i]*(fs[k-i]+2*(fs[k-1]-fs[k-i]));
        }
        for(int i=0;i<k;i++)g[i]+=f[i];
    }
}

void dfz(int u){
    work(u);
    for(int v:G[u]){
        if(vis[v])continue;
        rt=0,all=sz[v],calcsiz(v,u),calcsiz(rt,u),dfz(rt);
    }
}

signed main(void){

#ifndef ONLINE_JUDGE
    freopen("in.in","r",stdin);
#endif

    n=read(),k=read();
    for(int i=1;i<=n-1;i++){
        int u=read(),v=read();
        G[u].push_back(v),G[v].push_back(u);
    }
    all=n,calcsiz(1,0),calcsiz(rt,0),dfz(rt);
    cout<<ans<<endl;

    return 0;
}