CF771C Bear and Tree Jumps
云浅知处
2022-04-25 15:37:59
> 给定一棵 $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$，那么要求的就是：
$$
\sum_{b_u\neq b_v}\left\lceil\dfrac{d_u+d_v}{k}\right\rceil
$$
其中 $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)$，故
$$
\left\lceil\dfrac{d_u+d_v}{k}\right\rceil=\lfloor d_u/k\rfloor+\lfloor d_v/k\rfloor+\left\lceil\dfrac{(d_u\bmod k) +(d_v\bmod k)}{k}\right\rceil
$$
前一项可以直接算，考虑后一项怎么算。
设 $f_i$ 为当前子树内满足 $d_u\bmod k=i$ 的 $u$ 的个数，$g_i$ 为前面子树内 $\bmod k=i$ 的个数，那么后一项的和就是
$$
\sum_{i=0}^{k-1}\sum_{j=0}^{k-1}g_if_j\left\lceil\dfrac{i+j}{k}\right\rceil=\sum_{i=0}^{k-1}g_i\left(\sum_{j=1-i}^{k-i}f_j+2\sum_{j=k-i+1}^{k-1}f_j\right)
$$
维护 $f,g$ 以及 $f$ 的前缀和即可。$O(nk\log n)$。
感觉好像可以做到 $O((n+k)\log n)$ 之类的？
```cpp
#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;
}
```