SPOJ QTREE Query on a tree ——树链剖分 线段树

#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
#define maxn 20005

int T,n,fr[maxn],h[maxn],to[maxn],ne[maxn],w[maxn],en=0;
int mx[maxn<<3],fa[maxn],siz[maxn],son[maxn],pos[maxn],dep[maxn],tot=0;
int top[maxn],a[maxn],x,y,L,R,C,X,tofa[maxn];
char opt[10];

int Getint()
{
    int x=0,f=1; char ch=getchar();
    while (ch<‘0‘||ch>‘9‘) {if (ch==‘-‘) f=-1; ch=getchar();}
    while (ch>=‘0‘&&ch<=‘9‘) {x=x*10+ch-‘0‘; ch=getchar();}
    return x*f;
}

void add(int a,int b,int c)
{fr[en]=a;to[en]=b;ne[en]=h[a];w[en]=c;h[a]=en++;}

void rd()
{
	n=Getint();
	for(register int i=1;i<n;++i)
	{
		int a=Getint(),b=Getint(),c=Getint();
		add(a,b,c),add(b,a,c);
	}
}
void init()
{
	memset(fa,0,sizeof fa);
	memset(son,0,sizeof son);
	memset(h,-1,sizeof h);
	tot=en=0;
}

void dfs1(int o)
{
	siz[o]=1;
	for (register int i=h[o];i>=0;i=ne[i])
	if (fa[o]!=to[i]){
		fa[to[i]]=o;
		tofa[to[i]]=w[i];
		dep[to[i]]=dep[o]+1;
		dfs1(to[i]);
		siz[o]+=siz[to[i]];
		if (siz[son[o]]<siz[to[i]]) son[o]=to[i];
	}
}

void dfs2(int o,int tp,int ww)
{
//	printf("dfs2 _ %d %d %d son is %d\n",o,tp,ww,son[o]);
	top[o]=tp;pos[o]=++tot;a[pos[o]]=ww;
	if (son[o]==0) return;
//	for (register int i=h[o];i>=0;i=ne[i])
//		if (to[i]==son[o])
//			dfs2(son[o],tp,w[i]);
//	printf("%d to %d as height son\n",o,son[o]);
	dfs2(son[o],tp,tofa[son[o]]);
	for (register int i=h[o];i>=0;i=ne[i])
		if (fa[o]!=to[i]&&to[i]!=son[o])
			dfs2(to[i],to[i],w[i]);
}

void build(int o,int l,int r)
{
	int mid=l+r>>1;
	if (l==r){mx[o]=a[l];return;}
	build(o<<1,l,mid); build(o<<1|1,mid+1,r);
	mx[o]=max(mx[o<<1],mx[o<<1|1]);
}

void modify(int o,int l,int r)
{
	if (l==r) {mx[o]=C;return;}
	int mid=l+r>>1;
	if (X<=mid) modify(o<<1,l,mid);
	else modify(o<<1|1,mid+1,r);
	mx[o]=max(mx[o<<1],mx[o<<1|1]);
}

int query(int o,int l,int r)
{
	if (L<=l&&r<=R) return mx[o];
	int mid=l+r>>1;
	if (R<=mid) return query(o<<1,l,mid);
	if (L>mid) return query(o<<1|1,mid+1,r);
	else return max(query(o<<1,l,mid),query(o<<1|1,mid+1,r));
}

int ask(int x,int y)
{
	int ret=0;
	while (top[x]!=top[y])
	{
		if (dep[top[x]]<dep[top[y]]) swap(x,y);
		L=pos[top[x]];R=pos[x];
		ret=max(ret,query(1,1,n));
		x=fa[top[x]];
	}
	if (dep[x]<dep[y]) swap(x,y);
	if (x==y) return ret;
	L=pos[son[y]];R=pos[x];
	ret=max(ret,query(1,1,n));
	return ret;
}

int main()
{
	freopen("in.txt","r",stdin);
	freopen("wa.txt","w",stdout);
	T=Getint();
	while (T--)
	{
		init(),rd();
		dfs1(1);dfs2(1,1,0);
//		printf("n is %d\n",n);
		build(1,1,n);
		while (scanf("%s",opt)&&opt[0]!=‘D‘)
		{
			if (opt[0]==‘Q‘)
			{
				x=Getint();y=Getint();
				printf("%d\n",ask(x,y));
			}
			else
			{
				x=Getint(); y=Getint();
				L=fr[x*2-1];R=to[x*2-1];
				if (fa[L]==R) X=pos[L],C=y,modify(1,1,n);
				else X=pos[R],C=y,modify(1,1,n);
			}
		}
	}
	return 0;
}