南昌邀请赛 J. Distance on the tree

• J. Distance on the tree

题意：

给出n个节点的一颗树，每条边都有边权，m次查询，每次给出u,v,k，求u到v路径上边权小于等于k的边有多少条

思路：

对边权和k一起离散化，就把每个操作挂到对应的点上，跑一遍tarjan。向下搜索的时候边权w，就用树状数组维护加到第w个位置，回溯的时候删掉就行。假设有一个操作 (u,v,k)，那么假设搜索到u的时候，那直接就可以用树状数组求出根节点到u上小于等于k的边有a条(回溯的时候有删边，那么在树状数组中的数一定都是这条路径上的边权)，搜索到v的时候就可以直接求出根节点到v的路径上小于等于k的边有b条，然后此时还可以计算出fa=lca(u,v)，那么在fa这个点加个操作，等下回溯的时候算出根节点到fa的路径上小于等于x的边有c条，那最后就得到这个操作的答案a+b-2*c。

代码：

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int,int>pii;
typedef vector<int>vi;

#define rep(i,a,b) for(int i=(a);i<(b);i++)
#define fi first
#define se second
#define de(x) cout<<#x<<"="<<x<<endl;
#define dd(x) cout<<#x<<"="<<x<<" " ;
#define pb(x) push_back(x)
#define lowbit(a) ((a)&-(a))
#define per(i,a,b) for(int i=(b)-1;i>=(a);--i)
const int N=2e5+5;
int p[N],pre[N],ans[N],sum[N];
bool vis[N];
struct edge{
	int u,v,w;
}ed[N];
struct option{
	int u,v,k;
}opt[N];
struct node{
	int id,w,v;
	node(int tid,int tw,int tv){
		id=tid;
		w=tw;
		v=tv;
	}
};
map<int,int>mp,dis[N];
vector<pii>G[N],lca[N];
vector<node>op[N];
void add(int x,int y){
	for(int i=x;i<N;i+=lowbit(i)){
		sum[i]+=y;
	}
}
int getSum(int x){
	int res=0;
	for(int i=x;i>=1;i-=lowbit(i))
		res+=sum[i];
	return res;
}
int find_set(int x){
	if(pre[x]==x)return x;
	return pre[x]=find_set(pre[x]);
}
void tarjan(int u){
	vis[u]=true;
	for(auto it:G[u]){
		if(!vis[it.fi]){
			add(it.se,1);
			tarjan(it.fi);
			add(it.se,-1);
			pre[it.fi]=u;
		}
	}
	for(auto it:lca[u]){
		ans[it.se]-=2*getSum(it.fi);
	}
	for(auto it:op[u]){
		ans[it.id]+=getSum(it.w);
		if(vis[it.v]){
			int fa=find_set(it.v);
			lca[fa].pb(pii(it.w,it.id));	
		}
	}
}
int main()
{
	memset(vis,0,sizeof(vis));
	int n,m;
	scanf("%d%d",&n,&m);
	rep(i,1,n+1)
		pre[i]=i;	
	int len=n-1;
	rep(i,1,n){
		scanf("%d%d%d",&ed[i].u,&ed[i].v,&ed[i].w);
		p[i]=ed[i].w;
	}
	rep(i,1,m+1){
		scanf("%d%d%d",&opt[i].u,&opt[i].v,&opt[i].k);
		p[++len]=opt[i].k;
	}
	sort(p+1,p+len+1);
	int tot=unique(p+1,p+len+1)-p;
	for(int i=1;i<=len;++i){
		mp[p[i]]=lower_bound(p+1,p+tot+1,p[i])-p;
	}
	rep(i,1,n){
		G[ed[i].u].pb(pii(ed[i].v,mp[ed[i].w])); 
		G[ed[i].v].pb(pii(ed[i].u,mp[ed[i].w]));
	}
	rep(i,1,m+1){
		op[opt[i].u].pb(node(i,mp[opt[i].k],opt[i].v));
		op[opt[i].v].pb(node(i,mp[opt[i].k],opt[i].u));
	}
	tarjan(1);
	rep(i,1,m+1){
		printf("%d\n",ans[i]);
	}
	return 0;
}