Distance on the tree--树上主席树，树上差分(按边差分)

题目链接:

Distance on the tree - 题库 - 计蒜客

题意:

给定一棵n个节点带边权的树和m次询问，每次询问u到v的路径上权值小于等于k的边的个数。

第一行n,m。接下来n-1行u,v,w，表示节点u和v之间的边权为w，在接下来m行u,v,k，查询u到v的路径上权值小于等于k的边的个数。(2<=n<=1e5,1<=m<=1e5,1<=w,k<=1e9)

样例:

输入1

3 3
1 3 2
2 3 7
1 3 0
1 2 4
1 2 7

输出1

0
1
2

输入2

5 2
1 2 1000000000
1 3 1000000000
2 4 1000000000
3 5 1000000000
2 3 1000000000
4 5 1000000000

输出2

2
4

思路:

树上差分预备知识:

树上差分理解_枫茗、的博客-CSDN博客

代码:

离散化

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
const int inf=1e9+10;
const int N=1e5+10;
const int M=40*N;
int a[N],ca[N];//离散化的原始数组和去重后从大到小的数组 
int n,m;
int inv[N],num;

typedef struct d_n{
	int val,ki;
}d_n;
d_n dn[N];
bool cmp(d_n x,d_n y){
	return x.val<y.val;
}
void discretization(int a[],int inv[],d_n dn[]){
	for(int i=1;i<=n;i++){
		dn[i].val=a[i];
		dn[i].ki=i;
	}
	sort(dn+1,dn+n+1,cmp);
	int pre;
	for(int i=1;i<=n;i++){
		if(i==1||pre!=dn[i].val){
			++num;
			ca[num]=dn[i].val;
		}
		pre=dn[i].val;
		a[dn[i].ki]=num;
		inv[num]=dn[i].val;
	}
}


int root[N];
int lch[M],rch[M],tot;
int sum[M];
void cpy(int from,int to){
	lch[to]=lch[from];
	rch[to]=rch[from];
	sum[to]=sum[from]; 
}
void insert(int &u,int old,int L,int R,int x){
	u=++tot;
	cpy(old,u);
	sum[u]+=1;
	if(L==R){
		return ;
	}
		
	int mid=(L+R)>>1;
	if(x<=mid)
		insert(lch[u],lch[u],L,mid,x);
	else
		insert(rch[u],rch[u],mid+1,R,x);
}
int query(int s,int t,int ga,int L,int R,int l,int r){
	if(L==l&&R==r){
		return sum[t]+sum[s]-2*sum[ga]; 
		
	}
	int mid=(L+R)>>1;
	if(r<=mid){
		return query(lch[s],lch[t],lch[ga],L,mid,l,r);
	}
	else if(l<=mid){
		ll res=0;
		res+=query(lch[s],lch[t],lch[ga],L,mid,l,mid);
		res+=query(rch[s],rch[t],rch[ga],mid+1,R,mid+1,r);
		return res;
	}
	else
		return query(rch[s],rch[t],rch[ga],mid+1,R,l,r);
}

int h[N],to[2*N],ne[2*N],w[2*N],cnt;
int sz[N],dep[N],fa[N],son[N];
int top[N],id[N],dfn;

void add_edge(int u,int v,int ww){
	to[cnt]=v;
	ne[cnt]=h[u];
	w[cnt]=ww;
	h[u]=cnt++;
}
//轻重 
void dfs1(int u,int fat,int val){
	sz[u]=1;
	dep[u]=dep[fat]+1;
	fa[u]=fat;
	a[u]=val;
	for(int i=h[u];i!=-1;i=ne[i]){
		int v=to[i],ww=w[i];
		if(v!=fat){
			dfs1(v,u,ww);
			sz[u]+=sz[v];
			if(!son[u]||sz[son[u]]<sz[v]){
				son[u]=v;
			}
		}
	}
} 
//链剖分 
void dfs2(int u,int tp){
	id[u]=++dfn;
	insert(root[dfn],root[id[fa[u]]],1,num+1,a[u]);
	top[u]=tp;
	if(son[u])
		dfs2(son[u],tp);
	for(int i=h[u];i!=-1;i=ne[i]){
		if(to[i]!=fa[u]&&to[i]!=son[u]){
			dfs2(to[i],to[i]);
		}
	}
}
//求u和v的最近公共祖先 
int get_lca(int u,int v){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]){
			swap(u,v);
		}
		u=fa[top[u]];
	} 
	if(dep[u]>dep[v])
		swap(u,v);
	return u;
}

void init(){
	num=0;
	dfn=0;
	cnt=0;
	memset(h,-1,sizeof(h));
} 

int main(){
	init();
	scanf("%d%d",&n,&m);
	int u,v,ww,k,l,r,ga;
	for(int i=1;i<=n-1;i++){
		scanf("%d%d%d",&u,&v,&ww);
		add_edge(u,v,ww);
		add_edge(v,u,ww);
	}
	dfs1(1,0,0);
	//在insert之前要先进行离散化 
	discretization(a,inv,dn);
	dfs2(1,1);
	ca[num+1]=inf;
	while(m--){
		scanf("%d%d%d",&u,&v,&k);
		r=upper_bound(ca+1,ca+num+1,k)-ca-1;
		l=1;
		ga=get_lca(u,v);
		printf("%d\n",query(root[id[u]],root[id[v]],root[id[ga]],1,num+1,l,r));
	}
}

不离散化

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
const int inf=1e9+10;
const int N=1e5+10;
const int M=80*N;
int a[N];
int n,m;

int root[N];
int lch[M],rch[M],tot;
int sum[M];
void cpy(int from,int to){
	lch[to]=lch[from];
	rch[to]=rch[from];
	sum[to]=sum[from]; 
}
void insert(int &u,int old,int L,int R,int x){
	u=++tot;
	cpy(old,u);
	sum[u]+=1;
	if(L==R){
		return ;
	}
		
	int mid=(L+R)>>1;
	if(x<=mid)
		insert(lch[u],lch[u],L,mid,x);
	else
		insert(rch[u],rch[u],mid+1,R,x);
}
int query(int s,int t,int ga,int L,int R,int l,int r){
	if(L==l&&R==r){
		return sum[t]+sum[s]-2*sum[ga]; 
	}
	int mid=(L+R)>>1;
	if(r<=mid){
		return query(lch[s],lch[t],lch[ga],L,mid,l,r);
	}
	else if(l<=mid){
		ll res=0;
		res+=query(lch[s],lch[t],lch[ga],L,mid,l,mid);
		res+=query(rch[s],rch[t],rch[ga],mid+1,R,mid+1,r);
		return res;
	}
	else
		return query(rch[s],rch[t],rch[ga],mid+1,R,l,r);
}

int h[N],to[2*N],ne[2*N],w[2*N],cnt;
int sz[N],dep[N],fa[N],son[N];
int top[N],id[N],dfn;

void add_edge(int u,int v,int ww){
	to[cnt]=v;
	ne[cnt]=h[u];
	w[cnt]=ww;
	h[u]=cnt++;
}
//轻重 
void dfs1(int u,int fat,int val){
	sz[u]=1;
	dep[u]=dep[fat]+1;
	fa[u]=fat;
	a[u]=val;
	for(int i=h[u];i!=-1;i=ne[i]){
		int v=to[i],ww=w[i];
		if(v!=fat){
			dfs1(v,u,ww);
			sz[u]+=sz[v];
			if(!son[u]||sz[son[u]]<sz[v]){
				son[u]=v;
			}
		}
	}
} 
//链剖分 
void dfs2(int u,int tp){
	id[u]=++dfn;
	insert(root[dfn],root[id[fa[u]]],0,inf,a[u]);
	top[u]=tp;
	if(son[u])
		dfs2(son[u],tp);
	for(int i=h[u];i!=-1;i=ne[i]){
		if(to[i]!=fa[u]&&to[i]!=son[u]){
			dfs2(to[i],to[i]);
		}
	}
}
//求u和v的最近公共祖先 
int get_lca(int u,int v){
	while(top[u]!=top[v]){
		if(dep[top[u]]<dep[top[v]]){
			swap(u,v);
		}
		u=fa[top[u]];
	} 
	if(dep[u]>dep[v])
		swap(u,v);
	return u;
}

void init(){
	dfn=0;
	cnt=0;
	memset(h,-1,sizeof(h));
} 

int main(){
	init();
	scanf("%d%d",&n,&m);
	int u,v,ww,k,l,r,ga;
	for(int i=1;i<=n-1;i++){
		scanf("%d%d%d",&u,&v,&ww);
		add_edge(u,v,ww);
		add_edge(v,u,ww);
	}
	dfs1(1,0,0);
	dfs2(1,1);
	while(m--){
		scanf("%d%d%d",&u,&v,&k);
		ga=get_lca(u,v);
		printf("%d\n",query(root[id[u]],root[id[v]],root[id[ga]],0,inf,0,k));
	}
}