luogu P4178 Tree |点分治+树状数组
题目描述
给你一棵TREE,以及这棵树上边的距离.问有多少对点它们两者间的距离小于等于K
输入格式
N(n<=40000) 接下来n-1行边描述管道,按照题目中写的输入 接下来是k
输出格式
一行,有多少对点之间的距离小于等于k
如此简洁的题目在洛谷是真的少见啊
点分治统计,树状数组快速查找小于K的 O(nlogn*logn)
#include<cstdio>
#include<vector>
#include<cstring>
#include<iostream>
#include<algorithm>
const int N=5e5+10,inf=4e7;
using namespace std;
int n,K;
int nxt[N<<1],head[N],go[N<<1],w[N<<1],tot;
inline void add(int u,int v,int o){
nxt[++tot]=head[u]; head[u]=tot; go[tot]=v; w[tot]=o;
nxt[++tot]=head[v]; head[v]=tot; go[tot]=u; w[tot]=o;
}
int maxp[N],size[N],dis[N],rem[N];
int vis[N];
int sum,rt;
inline void getrt(int u,int fa){
size[u]=1; maxp[u]=0;
for(int i=head[u];i;i=nxt[i]){
int v=go[i];
if(v==fa)continue;
size[u]+=size[v];
maxp[u]=max(maxp[u],size[v]);
}
maxp[u]=max(maxp[u],sum-size[u]);
if(maxp[u]<maxp[rt])rt=u;
}
inline void getdis(int u,int fa){
if(dis[u]>K)return;
rem[++rem[0]]=dis[u];
for(int i=head[u];i;i=nxt[i]){
int v=go[i];
if(v==fa||vis[v])continue;
dis[v]=dis[u]+w[i];
getdis(v,u);
}
}
int c[N];
inline void update(int x,int y){
for(;x<=K;x+=x&(-x))c[x]+=y;
}
inline int ask(int x){
int ans=0;
for(;x;x-=x&(-x))ans+=c[x];
return ans;
}
int ans=0,q[N];
inline void calc(int u){
int p=0;
for(int i=head[u];i;i=nxt[i]){
int v=go[i];
if(vis[v])continue;
rem[0]=0; dis[v]=w[i];
getdis(v,u);
for(int j=rem[0];j;j--)ans+=ask(K-rem[j])+1;
for(int j=rem[0];j;j--)
update(rem[j],1),q[++p]=rem[j];
}
for(int i=1;i<=p;i++)update(q[i],ask(q[i]-1)-ask(q[i]));
}
inline void solve(int u){
vis[u]=1; calc(u);
for(int i=head[u];i;i=nxt[i]){
int v=go[i];
if(vis[v])continue;
sum=size[v];
maxp[rt=0]=inf;
getrt(v,0);
solve(rt);
}
}
signed main(){
cin>>n;
for(int i=1,u,v,o;i<n;i++){
scanf("%d%d%d",&u,&v,&o);
add(u,v,o);
}
cin>>K;
maxp[0]=sum=n;
getrt(1,0);
solve(rt);
cout<<ans<<endl;
}
不以物喜,不以己悲