题意:
给定一棵 n 个节点的树,每条边有边权,求出树上两点距离小于等于 k 的点对数量。
题解:
根点分治模板提很相似,只不过这个题目让你去维护小于等于k距离点的个数,这个时候我们还是要用到桶的思想,只不过每次查询查的时小于等于(k-某个值)的个数,对于这个操作,我们可以用一个树状数组来维护即可。
做的时候还是很多方面的小bug,比如说找重心的时候手残…等等
代码:
#include<bits/stdc++.h>
using namespace std;
const int maxn=1e5+10;
//int n,m;
int c[20005],k; //对应原数组和树状数组
int lowbit(int x){
return x&(-x);
}
void updata(int i,int k){ //在i位置加上k
while(i <= 20005){
c[i] += k;
i += lowbit(i);
}
}
int getsum(int i){ //求A[1 - i]的和 y-(x-1)
int res = 0;
while(i > 0){
res += c[i];
i -= lowbit(i);
}
return res;
}
struct E{
int to,next,w;
}edge[maxn];
int tot,head[maxn];
bool visited[maxn];
int maxp[maxn],sz[maxn],dis[maxn],temp[maxn];
int sum,rt,ans,n,cnt;
void add(int u,int v,int w){
edge[tot].to=v;
edge[tot].w=w;
edge[tot].next=head[u];
head[u]=tot++;
}
void getrt(int x,int fa){
sz[x]=1,maxp[x]=0;
for(int i=head[x];~i;i=edge[i].next){
int v=edge[i].to;
if(v==fa||visited[v]) continue;
getrt(v,x);
sz[x]+=sz[v];
if(sz[v]>maxp[x]) maxp[x]=sz[v];
}
maxp[x]=max(maxp[x],sum-sz[x]);
if(maxp[x]<maxp[rt]) rt=x;
}
void getdis(int x,int fa){
temp[cnt++]=dis[x];
//updata(dis[x],1);
if(dis[x]<=k) ans++;
for(int i=head[x];~i;i=edge[i].next){
int v=edge[i].to;
if(v==fa||visited[v]) continue;
dis[v]=dis[x]+edge[i].w;
getdis(v,x);
}
}
void sol(int x){
queue<int> que;
for(int i=head[x];~i;i=edge[i].next){
int v=edge[i].to;
if(visited[v]) continue;
dis[v]=edge[i].w;
cnt=0;
getdis(v,x);
// for(int j=0;j<cnt;j++) cout<<temp[j]<<" ";
// cout<<endl;
for(int j=0;j<cnt;j++){
ans+=getsum(k-temp[j]);
}
for(int j=0;j<cnt;j++){
que.push(temp[j]);
updata(temp[j],1);
//cout<<"sum "<<getsum(temp[j])<<endl;
}
}
while(!que.empty()){
updata(que.front(),-1);
que.pop();
}
}
void divide(int x){
visited[x]=true;
sol(x);
for(int i=head[x];~i;i=edge[i].next){
int v=edge[i].to;
if(visited[v]) continue;
maxp[rt=0]=sum=sz[v];
getrt(v,0);
getrt(rt,0);
divide(rt);
}
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
cout.tie(nullptr);
memset(head,-1,sizeof head);
cin>>n;
for(int i=1;i<n;i++){
int x,y,z;
cin>>x>>y>>z;
add(x,y,z);
add(y,x,z);
}
cin>>k;
maxp[0]=sum=n;
getrt(1,0);
//cout<<rt<<endl;
getrt(rt,0);
divide(rt);
cout<<ans<<endl;
}