Housewife Wind POJ - 2763.轻重链剖分

轻重链剖分裸题

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <memory.h>
#define mid ((l+r)>>1)
using namespace std;
const int maxn=2e5+5;
typedef long long ll;
struct E{
	int v,w,next;
};
E edge[maxn*2];

struct D{
	int u,v,w;
};
D ed[maxn];
int n,q,s,u,v,w;
int id[maxn];
int head[maxn];
int son[maxn];
int size[maxn];
int fa[maxn];
int top[maxn];
int dep[maxn];
int sum[maxn];
int tree[maxn<<2];
int lenid=0;
int tot=0;
void add(int u,int v,int w){
	edge[tot].v=v;
	edge[tot].w=w;
	edge[tot].next=head[u];
	head[u]=tot++;
}

void init(){
	memset(head,-1,sizeof(head));
	tot=lenid=0;
}

void dfs1(int u,int f,int d){
	dep[u]=d;
	size[u]=1;
	fa[u]=f;
	son[u]=0;
	for(int i=head[u];~i;i=edge[i].next){
		int v=edge[i].v;
		if(v!=f){
			dfs1(v,u,d+1);
			size[u]+=size[v];
			if(son[u]==0 || size[son[u]]<size[v]){
				son[u]=v;
			}
		}
	}
}

void dfs2(int u,int f){
	id[u]=++lenid;
//	cout<<u<<" "<<lenid<<endl;
	top[u]=f;
	if(son[u]){
		dfs2(son[u],f);
	}
	for(int i=head[u];~i;i=edge[i].next){
		int v=edge[i].v;
		if(v!=fa[u] && v!=son[u]){
			dfs2(v,v);
		}
	}
}

void pushup(int rt){
	tree[rt]=tree[rt<<1|1]+tree[rt<<1];
}

void build(int l,int r,int rt){
	if(l==r){
		tree[rt]=sum[l];
		return ;
	}
	build(l,mid,rt<<1);
	build(mid+1,r,rt<<1|1);
	pushup(rt);
}

void update(int x,int c,int l,int r,int rt){
	if(l==r){
		tree[rt]=c;
		return ;
	}
	if(x<=mid){
		update(x,c,l,mid,rt<<1);
	}else{
		update(x,c,mid+1,r,rt<<1|1);
	}
	pushup(rt);
}

int query(int a,int b,int l,int r,int rt){
	if(a<=l && b>=r){
		return tree[rt];
	}
	int ans=0;
	if(a<=mid){
		ans+=query(a,b,l,mid,rt<<1);
	}
	if(b>mid){
		ans+=query(a,b,mid+1,r,rt<<1|1);
	}
	pushup(rt);
	return ans;
}

int lca(int x,int y){
	int ans=0;
	while(top[x]!=top[y]){
		if(dep[top[x]]<dep[top[y]]){
			swap(x,y);
		}
		ans+=query(id[top[x]],id[x],1,n,1);
		x=fa[top[x]];
	}
	if(x==y){
		return ans;
	}
	if(dep[x]>dep[y]){
		swap(x,y);
	}
	ans+=query(id[son[x]],id[y],1,n,1);
	return ans;
}
/*
3 3 1
1 2 1
2 3 2
0 2
1 2 3
0 3
*/
int main(){
	while(scanf("%d%d%d",&n,&q,&s)!=EOF){
		init();
		for(int i=1;i<n;i++){
			scanf("%d%d%d",&u,&v,&w);
			ed[i].u=u;
			ed[i].v=v;
			ed[i].w=w;
			add(u,v,w);
			add(v,u,w);
		}
		dfs1(1,0,1);
		dfs2(1,0);
		for(int i=1;i<n;i++){
			if(dep[ed[i].u]<dep[ed[i].v]){
				swap(ed[i].u,ed[i].v);
			}
			sum[id[ed[i].u]]=ed[i].w;
		}
//		for(int i=1;i<=n;i++){
//			cout<<sum[i]<<" ";
//		}cout<<endl;
		build(1,n,1);
		while(q--){
			int pos;
			scanf("%d",&pos);
			if(pos==0){
				int p;
				scanf("%d",&p);
				printf("%d\n",lca(s,p));
				s=p;
			}else{
				int l,w;
				scanf("%d%d",&l,&w);
				int us=ed[l].u;
				update(id[us],w,1,n,1);
			}
		}
	}
	return 0;
}