(南昌邀请网络赛) J.Distance on the tree (树链剖分+主席树)

传送门

题意：给一颗树，m次查询ui->vi这条链中边权小于等于ki的边数。

解：树链剖分+主席树思路很清晰，可以解决；（又因为数组开小了而狂T不止

#include<bits/stdc++.h>
#define il inline
#define pb push_back
#define ms(_data,v) memset(_data,v,sizeof(_data))
#define SZ(a) int((a).size())
#define mid ((l+r)>>1)
using namespace std;
typedef long long ll;
const ll inf=0x3f3f3f3f;
const int N=2e5+5; 
template <typename _Tp> il void read(_Tp&x) {
	char ch;bool flag=0;x=0;
	while(ch=getchar(),!isdigit(ch)) if(ch=='-')flag=1;
	while(isdigit(ch)) x=x*10+ch-'0',ch=getchar();
	if(flag) x=-x;
}
//il int Add(ll &x,ll y) {return x=x+y>=mod?x+y-mod:x+y;}
//il int Mul(ll &x,ll y) {return x=x*y>=mod?x*y%mod:x*y;}
struct node{int to,w;};
vector<node> G[N];
int dep[N],fa[N],son[N],sz[N],ed[N],ned[N],id[N],bel[N],cnt;
il void dfs1(int x,int ff){
	dep[x]=dep[ff]+1;
	fa[x]=ff,sz[x]=1;
	int mx=-1,to,w,len=SZ(G[x]);
	node tp;
	for(int i=0;i<len;++i){
		tp=G[x][i],to=tp.to,w=tp.w;
		if(to==ff) continue;
		ed[to]=w;
		dfs1(to,x);
		sz[x]+=sz[to];
		if(sz[to]>mx) son[x]=to,mx=sz[to];
	}
} 
il void dfs2(int x,int topx){
	id[x]=++cnt,ned[cnt]=ed[x];
	bel[x]=topx;
	if(!son[x]) return;
	dfs2(son[x],topx);
	int to,len=SZ(G[x]);
	node tp;
	for(int i=0;i<len;++i){
		tp=G[x][i],to=tp.to;
		if(to==son[x] || to==fa[x]) continue;
		dfs2(to,to);
	}
} 
int n,m,b[N],tot=0,rt[N*20],ls[N*20],rs[N*20],s[N*20],all=0;
il void build(int &rt,int l,int r){
	rt=++tot,s[rt]=0;
	if(l==r) return;
	build(ls[rt],l,mid),build(rs[rt],mid+1,r);
} 
il void update(int &rt,int l,int r,int last,int p){
	rt=++tot,ls[rt]=ls[last],rs[rt]=rs[last];
	s[rt]=s[last]+1;
	if(l==r) return;
	if(p<=mid) update(ls[rt],l,mid,ls[last],p);
	else update(rs[rt],mid+1,r,rs[last],p);
}
il int query(int ss,int ee,int l,int r,int k){
	if(l==r) return s[ee]-s[ss];
	if(k>mid) return s[ls[ee]]-s[ls[ss]]+query(rs[ss],rs[ee],mid+1,r,k);
	else if(k==mid) return s[ls[ee]]-s[ls[ss]];
	else return query(ls[ss],ls[ee],l,mid,k);	
}
il int r_query(int x,int y,int k){
//	cout<<"r_query "<<x<<" "<<y<<" "<<k<<endl;
//	cout<<bel[x]<<" "<<bel[y]<<endl;
	int res=0;
	while(bel[x]!=bel[y]){
		if(dep[bel[x]]<dep[bel[y]]) swap(x,y);
		res+=query(rt[id[bel[x]]-1],rt[id[x]],1,all,k);
		x=fa[bel[x]];
	}
	if(dep[x]>dep[y]) swap(x,y);
	res+=query(rt[id[x]],rt[id[y]],1,all,k);
	return res;
}
struct Q{
	int l,r,k;
}qu[N];
int main(){
//	std::ios::sync_with_stdio(0);cin.tie(0);
	read(n),read(m);
	int x,y,w,tt=0;
	for(int i=1;i<=n-1;++i){
		read(x),read(y),read(w);
		G[x].pb({y,w});
		G[y].pb({x,w});
		b[++tt]=w;
	}
	for(int i=1;i<=m;++i){
		read(x),read(y),read(w);
		qu[i]={x,y,w};
		b[++tt]=w;
	}
	dfs1(1,0),dfs2(1,1);
	b[++tt]=0;
	sort(b+1,b+tt+1);
	all=unique(b+1,b+tt+1)-(b+1);
	build(rt[0],1,all);
	for(int i=1;i<=n;++i){
		ned[i]=lower_bound(b+1,b+all+1,ned[i])-b;
		update(rt[i],1,all,rt[i-1],ned[i]);
	}
	for(int i=1;i<=m;++i){
		qu[i].k=lower_bound(b+1,b+all+1,qu[i].k)-b;
		printf("%d\n",r_query(qu[i].l,qu[i].r,qu[i].k));
	}
	return 0;
}