Query on a tree II

终于用树链剖分水掉了这题,LCA这种东西,树剖能做到O(N)预处理,O(logN)查询,所以果断不用倍增(每次都写错TAT)。

数据貌似很坑爹的样子,第一次WA了,目测有o-th的情况,于是加了个特判,过了,0.06S

好纠结以后LCA用什么做啊。

#include<iostream>
#include<cstdio>
#include<cstring>
using namespace std;
const int N=10000+5;
struct Edge{int to,next,v;}e[N*2];
int head[N],cnt;
void ins(int u,int v,int w){
	e[++cnt].to=v;e[cnt].next=head[u];head[u]=cnt;e[cnt].v=w;
}
int fa[N],dep[N],siz[N],son[N],top[N],pos[N],rank[N],sz;
int d[N];
void dfs1(int u){
	son[u]=0;siz[u]=1;
	for(int i=head[u];i;i=e[i].next){
		int v=e[i].to;
		if(v==fa[u])continue;
		dep[v]=dep[u]+1;fa[v]=u;d[v]=d[u]+e[i].v;
		dfs1(v);
		siz[u]+=siz[v];
		if(siz[v]>siz[son[u]])son[u]=v;
	}
}
void dfs2(int u,int tp){
	pos[u]=++sz;top[u]=tp;rank[sz]=u;
	if(son[u])dfs2(son[u],tp);
	for(int i=head[u];i;i=e[i].next){
		int v=e[i].to;
		if(v!=fa[u]&&v!=son[u])dfs2(v,v);
	}
}
int lca(int u,int v){
	while(top[u]!=top[v]){
		if(dep[top[u]]>dep[top[v]])u=fa[top[u]];
		else v=fa[top[v]];
	}
	if(dep[u]<dep[v])return u;
	return v;
}
int dist(int u,int anc){
	return dep[u]-dep[anc]+1;
}
int kth(int from,int k){
	if(k>dist(from,top[from]))return kth(fa[top[from]],k-dist(from,top[from]));
	else return rank[pos[from]-k+1];
}
int find(int u,int v,int k){
	if(!k)return 0;
	int w=lca(u,v);
	if(k<=dist(u,w))return kth(u,k);
	else return kth(v,dep[u]+dep[v]-2*dep[w]+1-k+1);
}
int main(){
	int T;scanf("%d",&T);
	while(T--){
		int n;scanf("%d",&n);
		int u,v,w;
		memset(head,0,sizeof(head));cnt=0;sz=0;
		for(int i=1;i<n;i++){
			scanf("%d%d%d",&u,&v,&w);
			ins(u,v,w);ins(v,u,w);
		}
		dfs1(1);dfs2(1,1);
		char opt[10];int a,b,k;
		while(true){
			scanf("%s",opt);
			if(opt[1]=='O')break;
			scanf("%d%d",&a,&b); 
			if(opt[0]=='D')printf("%d\n",d[a]+d[b]-2*d[lca(a,b)]);
			else{
				scanf("%d",&k);
				printf("%d\n",find(a,b,k));
			}
		}
	}
	return 0;
}


你可能感兴趣的:(Query on a tree II)