BZOJ 3991: [SDOI2015]寻宝游戏 树链的并+set

#include <set>    
#include <cstdio>  
#include <algorithm> 
#define N 200003 
#define ll long long  
#define setIO(s) freopen(s".in","r",stdin) , freopen(s".out","w",stdout)     
using namespace std;   
set<int>s;   
set<int>::iterator it;          
ll dis[N],now;
int n,m,edges,tim; 
int hd[N],to[N<<1],nex[N<<1],val[N<<1]; 
int top[N],fa[N],dep[N],dfn[N],size[N],son[N],re[N];        
int vis[N]; 
void add(int u,int v,int c) {
	nex[++edges]=hd[u],hd[u]=edges,to[edges]=v,val[edges]=c; 
} 
void dfs1(int u,int ff) {
	fa[u]=ff,size[u]=1,dfn[u]=++tim,re[tim]=u;          
	for(int i=hd[u];i;i=nex[i]) 
		if(to[i]!=ff) {
			dep[to[i]]=dep[u]+1,dis[to[i]]=dis[u]+1ll*val[i]; 
			dfs1(to[i],u), size[u]+=size[to[i]]; 
			if(size[to[i]]>size[son[u]]) son[u]=to[i]; 
		}
}
void dfs2(int u,int tp) {
	top[u]=tp; 
	if(son[u]) dfs2(son[u],tp); 
	for(int i=hd[u];i;i=nex[i]) 
		if(to[i]!=fa[u]&&to[i]!=son[u]) 
			dfs2(to[i],to[i]);             
} 
int LCA(int x,int y) {
	while(top[x]!=top[y]) 
		dep[top[x]]>dep[top[y]]?x=fa[top[x]]:y=fa[top[y]]; 
	return dep[x]<dep[y]?x:y;    
}    
int main() { 
	int i,j; 
	// setIO("input");     
	scanf("%d%d",&n,&m);  
	for(i=1;i<n;++i) {
		int x,y,z; 
		scanf("%d%d%d",&x,&y,&z),add(x,y,z),add(y,x,z); 
	}   
	dfs1(1,0), dfs2(1,1);  
	for(i=1;i<=m;++i) {
		int x; 
		scanf("%d",&x);                               
		if(!vis[x]) {
			int y=0,z=0; 
			it=s.upper_bound(dfn[x]);    
			if(it!=s.end()) z=(*it); 
			if(it!=s.begin()) it--, y=(*it);    
			if(z&&y) now+=dis[LCA(re[z],re[y])]; 
			if(z) now-=dis[LCA(re[z],x)]; 
			if(y) now-=dis[LCA(x,re[y])];         
			s.insert(dfn[x]),vis[x]=1,now+=dis[x];      
		}
		else { 
			now-=dis[x], vis[x]=0, s.erase(dfn[x]);      
			int z=0,y=0;  
			it=s.upper_bound(dfn[x]); 
			if(it!=s.end()) z=(*it);       
			if(it!=s.begin()) it--, y=(*it);    
			if(y) now+=dis[LCA(re[y],x)];                 
			if(z) now+=dis[LCA(x,re[z])];     
			if(y&&z) now-=dis[LCA(re[y],re[z])];     
		}   
		if(s.size()<2) printf("0\n"); 
		else {
			int a,b; 
			it=s.begin(), a=(*it); 
			it=s.end(), it--, b=(*it);  
			printf("%lld\n",(now-dis[LCA(re[a],re[b])])<<1);             
		}        
	}
	return 0; 
}