一、主要思想
很容易理解，我们将一个树以一个节点分割成若干个子树。对于这个节点，我们以一些方式统计和改变答案，然后不断地向子树递归。
那应该选择哪个节点呢？显然是重心。树的重心有一个性质：所有子树的大小小于等于当前树的大小的二分之一。也就是说，这保证了递归层数 \(log_2\) 的时间复杂度。
二、实现
#include<bits/stdc++.h>
using namespace std;
const int N=40010;
int n,cnt,first[N],k,root,p,ans;
int siz[N],rootsiz,dis[N];
bool vis[N];
struct edeg
{
int to,v,next;
}e[N*2];
void addedge(int x,int y,int z)
{
e[++cnt].to=y;
e[cnt].v=z;
e[cnt].next=first[x];
first[x]=cnt;
}
void findroot(int now,int fa,int all)
{
int nowmax=0;
siz[now]=1;
for(int i=first[now];i;i=e[i].next)
{
int to=e[i].to;
if(to==fa||vis[to])continue;
findroot(to,now,all);
siz[now]+=siz[to];
if(siz[to]>nowmax)
nowmax=siz[to];
}
if(all-siz[now]>nowmax)
nowmax=all-siz[now];
if(nowmax<rootsiz)
{
root=now;
rootsiz=nowmax;
}
}
void getdis(int now,int fa,int dv)
{
siz[now]=1;
dis[++p]=dv;
for(int i=first[now];i;i=e[i].next)
{
int to=e[i].to;
if(to==fa||vis[to])continue;
getdis(to,now,dv+e[i].v);
siz[now]+=siz[to];
}
}
void cc(int x)//计算答案
{
sort(dis+1,dis+p+1);
for(int i=p;i>=1;i--)
{
int r=upper_bound(dis+1,dis+i,k-dis[i])-(dis+1);
ans+=r*x;
}
}
void calc(int now,int fa,int all)
{
rootsiz=100000;
findroot(now,fa,all);
p=0;
vis[now]=1;
getdis(root,0,0);
cc(1);
//此处使用容斥原理。先算所有答案，再减去同一颗子树内的答案，也可以用树形背包之类的。
for(int i=first[root];i;i=e[i].next)
{
int to=e[i].to;
if(to==fa||vis[to])continue;
p=0;
getdis(to,root,e[i].v);//计算距离
cc(-1);
}
for(int i=first[root];i;i=e[i].next)
{
int to=e[i].to;
if(to==fa||vis[to])continue;
calc(to,root,all);
}
}
int main()
{
cin>>n;
for(int i=1;i<n;i++)
{
int x,y,z;
cin>>x>>y>>z;
addedge(x,y,z);
addedge(y,x,z);
}
cin>>k;
calc(1,0,n);
cout<<ans;
return 0;
}
这是求树上点对有几个小于等于 \(k\) 的。
注意，只要可以通过一些方式计算当前点贡献的答案的，几乎都可以用点分治。
三、例题选讲
1. \(\color{purple}{P3714 [BJOI2017] 树的难题}\)
毒瘤题，希望再也不要遇到了。
首先肯定点分治，难点是怎么计算答案。应该把这条边所连的所有边按颜色排序，然后一边跑一边用两颗线段树维护，一个同色，一个异色，问题解决了。
#include<algorithm>
#include<list>
#include<vector>
#include<iostream>
#pragma GCC optimize("Ofast","O2")
using namespace std;
const int N=200010,inf=2000000009;
int n,m,sum[N<<2][2],tag[N<<2][2];
int siz[N],dis[N],dis2[N],rootsiz=N;
int root,p,ans=-inf,le,ri,c[N];
bool vis[N];
int read()
{
int x,c,op=1;
while(c=getchar(),c<'0'||c>'9')
if(c=='-')op=-op;
x=c^48;
while(c=getchar(),c>='0'&&c<='9')
x=(x<<3)+(x<<1)+(c^48);
return x*op;
}
struct dot
{
int v,col;
};
vector<dot> e[N];
bool cmp(dot a,dot b)
{
return a.col<b.col;
}
void build(int p,int l,int r,int num)
{
tag[p][num]=0;
if(l==r){sum[p][num]=-inf;return ;}
int mid=(l+r)>>1;
build(p*2,l,mid,num);
build(p*2+1,mid+1,r,num);
sum[p][num]=-inf;
}
void f(int p,int l,int r,int k,int num)
{
tag[p][num]=k;
sum[p][num]=k;
}
void push_down(int p,int l,int r,int num)
{
int mid=(l+r)>>1;
if(tag[p][num]==0)return;
f(p*2,l,mid,tag[p][num],num);
f(p*2+1,mid+1,r,tag[p][num],num);
tag[p][num]=0;
}
void update(int nl,int nr,int l,int r,int p,int k,int num)
{
if(nl<=l&&r<=nr)
{
sum[p][num]=max(sum[p][num],k);
return ;
}
push_down(p,l,r,num);
int mid=(l+r)>>1;
if(nl<=mid)update(nl,nr,l,mid,p*2,k,num);
if(nr>mid) update(nl,nr,mid+1,r,p*2+1,k,num);
sum[p][num]=max(sum[p*2][num],sum[p*2+1][num]);
}
void clean(int nl,int nr,int l,int r,int p,int k,int num)
{
if(nl<=l&&r<=nr)
{
sum[p][num]=-inf;
tag[p][num]=-inf;
return ;
}
push_down(p,l,r,num);
int mid=(l+r)>>1;
if(nl<=mid)update(nl,nr,l,mid,p*2,k,num);
if(nr>mid) update(nl,nr,mid+1,r,p*2+1,k,num);
sum[p][num]=max(sum[p*2][num],sum[p*2+1][num]);
}
int query(int x,int y,int l,int r,int p,int num)
{
if(x==0)x=1;
if(y==0)return -inf;
int res=-inf;
if(x<=l&&r<=y)return sum[p][num];
int mid=(l+r)>>1;
push_down(p,l,r,num);
if(x<=mid)res=max(res,query(x,y,l,mid,p*2,num));
if(y>mid)res=max(res,query(x,y,mid+1,r,p*2+1,num));
return res;
}
void findroot(int now,int fa,int all)
{
int nowmax=0;
siz[now]=1;
for(int i=0;i<e[now].size();i++)
{
int to=e[now][i].v;
if(to==fa||vis[to])continue;
findroot(to,now,all);
siz[now]+=siz[to];
if(siz[to]>nowmax)
nowmax=siz[to];
}
if(all-siz[now]>nowmax)
nowmax=all-siz[now];
if(nowmax<rootsiz)
{
root=now;
rootsiz=nowmax;
}
}
void getdis(int now,int fa,int dv,int dl,int lastcol)
{
siz[now]=1;
dis[++p]=dv;
dis2[p]=dl;
for(int i=0;i<e[now].size();i++)
{
int to=e[now][i].v;
int nowcol=e[now][i].col;
if(to==fa||vis[to])continue;
if(e[now][i].col==lastcol)
getdis(to,now,dv,dl+1,nowcol);
else getdis(to,now,dv+c[nowcol],dl+1,nowcol);
siz[now]+=siz[to];
}
}
void calc(int now,int fa,int all)
{
rootsiz=N;
findroot(now,fa,all);
p=0;
int l=0,lcol=0,top=1;
dis[1]=0,dis2[1]=0;
vis[root]=1;
for(int i=0;i<e[root].size();i++)
{
int to=e[root][i].v;
if(to==fa||vis[to])continue;
getdis(to,root,c[e[root][i].col],1,e[root][i].col);
if(e[root][i].col==lcol)
{
for(int j=l+1;j<=p;j++)
{
int lask,rask;
if(dis2[j]<le)
lask=le-dis2[j],rask=ri-dis2[j];
else if(dis2[j]>=le&&dis2[j]<=ri)
lask=0,rask=ri-dis2[j],ans=max(ans,dis[j]);
else lask=0,rask=0;
int t1=query(lask+1,rask+1,1,n+1,1,0),t2=query(lask+1,rask+1,1,n+1,1,1);
if(t1!=-inf)ans=max(ans,t1-c[e[root][i].col]+dis[j]);
if(t2!=-inf)ans=max(ans,t2+dis[j]);
}
for(int j=l+1;j<=p;j++)
{
update(dis2[j]+1,dis2[j]+1,1,n+1,1,dis[j],0);
}
}
else
{
for(int j=top;j<=l;j++)
update(dis2[j]+1,dis2[j]+1,1,n+1,1,dis[j],1);
for(int j=l+1;j<=p;j++)
{
int lask,rask;
if(dis2[j]<le)
lask=le-dis2[j],rask=ri-dis2[j];
else if(dis2[j]>=le&&dis2[j]<=ri)
lask=0,rask=ri-dis2[j],ans=max(ans,dis[j]);
else lask=0,rask=0;
int t1=query(lask+1,rask+1,1,n+1,1,1);
if(t1!=-inf)ans=max(ans,t1+dis[j]);
}
clean(1,n+1,1,n+1,1,-inf,0);
for(int j=l+1;j<=p;j++)
update(dis2[j]+1,dis2[j]+1,1,n+1,1,dis[j],0);
top=l+1;
}
lcol=e[root][i].col;
l=p;
}
clean(1,n+1,1,n+1,1,-inf,1);
clean(1,n+1,1,n+1,1,-inf,0);
int tysroot=root;
for(int i=0;i<e[tysroot].size();i++)
{
int to=e[tysroot][i].v;
if(to==fa||vis[to])continue;
calc(to,tysroot,siz[to]);
}
}
signed main()
{
n=read();
m=read();
le=read();
ri=read();
for(int i=1;i<=m;i++)
c[i]=read();
for(int i=1;i<n;i++)
{
int x=read(),y=read(),z=read();
dot ax,bx;
ax.v=y,ax.col=z;
bx.v=x,bx.col=z;
e[x].push_back(ax);
e[y].push_back(bx);
}
build(1,1,n+1,0);
build(1,1,n+1,1);
for(int i=1;i<=n;i++)
{
stable_sort(e[i].begin(),e[i].end(),cmp);
}
calc(1,0,n);
printf("%d",ans);
return 0;
}