BJOI2017 树的难题 点分治+线段树

原题：https://loj.ac/submission/420518

题解:要找到长度L-R的染色的最大权值。长度L-R肯定要用点分治，通过分析可以得出对于每一个子树，只要考虑当前分治中心到每颗子树的颜色是不是一样的就行了。即相同的要减去1个。我们考虑用两颗线段树维护。线段树的坐标表示长度。对于每一组询问，要从[max(0,L-num),y-num]中找到最大值。为了方便起见，把所有长度都+1。由于要在点分治加线段树，所以常数很大，能用的优化尽量用。。。代码量挺大的。在洛谷上过不去

#include<bits/stdc++.h>
#pragma comment(linker, "/STACK:102400000,102400000")
#define inf 0x3f3f3f3f
#define reg register
using namespace std;

const int N=2e5+10;

struct E{int to,w,nxt;}data[N<<1];//w为颜色种类 

int n,m,L,R,len=1,rt,top,ans;
int clr[N],h[N],vis[N],f[N],sz[N];

inline int rd(){
	int x=0;int f=1;char s=getchar();
	while(!isdigit(s)) f=(s=='-'?-1:f),s=getchar();
	while(isdigit(s)) x=(x<<1)+(x<<3)+s-'0',s=getchar();
	return x*f;
}
inline void ins(int x,int y,int w){
	data[++len].to=y;data[len].w=w;data[len].nxt=h[x];h[x]=len;
	data[++len].to=x;data[len].w=w;data[len].nxt=h[y];h[y]=len;
} 

struct segmentTree{
	int tt[N<<2];bool tag[N<<2];
	#define ls now<<1
	#define rs now<<1|1
	#define mid (l+r)/2
	#define lson ls,l,mid
	#define rson rs,mid+1,r 
	inline void clear(){
		tag[1]=1;tt[1]=-inf;return ;
	}
	inline void pushdown(int now){
		if(tag[now]){
			tag[ls]=tag[rs]=1;tt[ls]=tt[rs]=-inf;tag[now]=0;
		} 
	}
	inline void upd(int now){
		tt[now]=max(tt[ls],tt[rs]);
	}
	inline void change(int now,int l,int r,int x,int c){
		if(l==x && r==x) {
			tt[now]=max(tt[now],c);return ;
		}
		tt[now]=max(tt[now],c);//常数优化 
		pushdown(now);
		if(x<=mid) change(lson,x,c);
		if(mid<x)  change(rson,x,c);
	//	upd(now);
	}
	inline int query(int now,int l,int r,int x,int y){
		if(x<=l && r<=y){return tt[now];}
		if(tag[now]) return -inf;
		pushdown(now);
		int ans=-inf;
		if(x<=mid) ans=max(ans,query(lson,x,y));
		if(mid<y)  ans=max(ans,query(rson,x,y));
		return ans;
	}
	inline void change(int p,int c){
		change(1,1,R+1,p+1,c); return ;
	} 
	inline int  query(int x,int y){
		if(x>y||y<0) return -inf;
		return query(1,1,R+1,max(1,x+1),y+1);
	}
	#undef ls
	#undef rs
	#undef lson
	#undef rson
	#undef mid 
}A,B;//A表示不同颜色的，B相同颜色来自不同的 
struct node{
	int dis,num,from,c;//从哪来的边 
}a[N];
bool cmp(node x,node y){//按颜色为第一关键字,来的边为第二关键字 
	return x.c==y.c?x.from<y.from:x.c<y.c;
} 
void getroot(int x,int fa){
	sz[x]=1;f[x]=0;
	for(int i=h[x];i;i=data[i].nxt){
		int y=data[i].to;if(y==fa || vis[y]) continue;
		getroot(y,x);
		sz[x]+=sz[y];
		f[x]=max(f[x],sz[y]);
	}
	f[x]=max(f[x],n-sz[x]); 
	if(f[rt]>f[x]) rt=x;
}
// dis距离 num长度 fclr 上一个的颜色 
void dfs(int x,int fa,int dis,int num,int fclr,int from,int c){
	if(num>R) return ;
	a[++top]=(node){dis,num,from,c};
	for(reg int i=h[x];i;i=data[i].nxt){
		int y=data[i].to;if(vis[y] || y==fa) continue;
		int w=data[i].w;
		if(fclr != w) dfs(y,x,dis+clr[w],num+1,w,from,c);
		else  dfs(y,x,dis,num+1,w,from,c);
	} 
}

void work(int x){
	a[top=1]=(node){0,0,0,0};
	for(int i=h[x];i;i=data[i].nxt){
		int y=data[i].to;if(vis[y]) continue;
		int w=data[i].w;
		dfs(y,x,clr[w],1,w,i,w);
	}
	sort(a+1,a+top+1,cmp);
	A.clear();B.clear();
	for(reg int l=1,r;l<=top;l=r+1){
		for(r=l;r<top &&a[l].c==a[r+1].c;r++);//颜色相同的分块
		for(reg int x=l,y;x<=r;x=y+1){//来自同一条边 
			for(y=x;y<r && a[x].from==a[y+1].from;y++); 
			if(x!=l) {
				for(int i=x;i<=y;i++) 
					ans=max(ans,a[i].dis+B.query(L-a[i].num,R-a[i].num)-clr[a[i].c]);
			}
			for(int i=x;i<=y;i++) B.change(a[i].num,a[i].dis);
		} 
		B.clear();
		if(l!=1) for(int i=l;i<=r;i++) ans=max(ans,a[i].dis+A.query(L-a[i].num,R-a[i].num));
		for(int i=l;i<=r;i++) A.change(a[i].num,a[i].dis);
	}
}
void sol(int x){
	vis[x]=1;work(x);
	for(int i=h[x];i;i=data[i].nxt){
		int y=data[i].to;if(vis[y]) continue;
		n=sz[y];rt=0;getroot(y,x);
		sol(rt); 
	}		
}
int main(){	
//	freopen("journey1.in","r",stdin);	
//	freopen("tree.in","r",stdin); 
	n=rd();m=rd();L=rd();R=rd(); 
	f[0]=inf;rt=0;
	for(int i=1;i<=m;i++) clr[i]=rd();
	for(int i=1,x,y,w;i<n;i++) {
		x=rd();y=rd();w=rd();
		ins(x,y,w);
	}
	ans=-2e9;
	getroot(1,0);// 
	sol(rt);
	printf("%d\n",ans);
	return 0;
}