【JZOJ5898】距离统计

Description

给你一颗N个点有边权的树，多次询问到一个点距离第k小为多少。

Solution

考虑点分治，把点分树建出来，对于一个询问点，直接往它的点分树父亲上跳。
求第k小考虑二分，那么相当于求权值和小于等于一个值的路径条数，那么计算跨越询问点的每一个点分树上的祖先的路径数即可。这里注意一下要减去重复的部分。

Code

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
#include<vector>
#define fo(i,j,k) for(int i=j;i<=k;++i)
#define fd(i,j,k) for(int i=j;i>=k;--i)
#define rep(i,x) for(int i=ls[x];i;i=nx[i])
using namespace std;
const int N=5e4+10,M=1e5+10,inf=1e9;
typedef long long ll;
int to[M],nx[M],ls[N],vl[M],num=0;
void link(int u,int v,int w){
	to[++num]=v,nx[num]=ls[u],ls[u]=num;
	vl[num]=w;
}
int sz[N],f[N],dep[N];
int fa[N][20],d[N][20];
vector<int> g[N],h[N];
int rt;
bool vis[N];
void getsz(int x,int fr,int tt){
	sz[x]=1,f[x]=0;
	rep(i,x){
		int v=to[i];
		if(v==fr || vis[v]) continue;
		getsz(v,x,tt),sz[x]+=sz[v];
		f[x]=max(f[x],sz[v]);
	}
	f[x]=max(f[x],tt-sz[x]);
	if(f[rt]>f[x]) rt=x;
}
void get(int x,int fr,int t,int now){
	sz[x]=1;
	g[now].push_back(d[x][dep[now]]);
	if(dep[now]>1) h[now].push_back(d[x][dep[now]-1]);
	rep(i,x){
		int v=to[i];
		if(v==fr || vis[v]) continue;
		d[v][dep[now]]=d[x][dep[now]]+vl[i];
		get(v,x,t,now),sz[x]+=sz[v];
	}
}
void dfs(int x,int fr){
	vis[x]=1,dep[x]=dep[fr]+1;
	fa[x][dep[x]]=x;
	fo(i,1,dep[fr]) fa[x][i]=fa[fr][i];
	g[x].push_back(0);
	if(fr) h[x].push_back(d[x][dep[fr]]);
	rep(i,x){
		int v=to[i];
		if(v==fr || vis[v]) continue;
		d[v][dep[x]]=vl[i],get(v,x,v,x);
	}
	rep(i,x){
		int v=to[i];
		if(v==fr || vis[v]) continue;
		rt=0,getsz(v,x,sz[v]);
		dfs(rt,x);
	}
}
ll calc(int x,int t){
	ll ans=0;
	fd(i,dep[x],1){
		int now=fa[x][i];
		ans+=upper_bound(g[now].begin(),g[now].end(),t-d[x][i])-g[now].begin();
		if(i<dep[x]){
			int la=fa[x][i+1];
			ans-=upper_bound(h[la].begin(),h[la].end(),t-d[x][i])-h[la].begin();
		}
	}
	return ans;
}
int main()
{
	freopen("tree.in","r",stdin);
	freopen("tree.out","w",stdout);
	int n,q;
	scanf("%d %d",&n,&q);
	fo(i,2,n){
		int u,v,w;
		scanf("%d %d %d",&u,&v,&w);
		link(u,v,w),link(v,u,w);
	}
	f[0]=n+1;
	rt=0,getsz(1,0,n);
	dep[rt]=1,dfs(rt,0);
	fo(i,1,n){
		int o=h[i].size();
		fo(j,0,o-1) printf("%d ",h[i][j]);
		printf("\n");
		sort(g[i].begin(),g[i].end()),sort(h[i].begin(),h[i].end());
		g[i].push_back(inf),h[i].push_back(inf);
	}
	while(q--){
		int u,k;
		scanf("%d %d",&u,&k),k++;
		int l=1,r=5e7;
		while(l+1<r){
			int mid=l+r>>1;
			calc(u,mid)>=k?r=mid:l=mid;
		}
		calc(u,l)>=k?r=l:0;
		printf("%d\n",r);
	}
}