BZOJ4860: [Beijing2017]树的难题

题面

对于每个重心，每个子树可以求出一个数组，表示每个深度的答案。先同种颜色合并，再不同颜色合并。合并可以按高度从小到大排序，然后跑单调队列，复杂度$O(n\log n)$。我比较脑抽，写的是从大到小排序，然后跑线段树，复杂度$O(n\log^2 n)$。

#include<bits/stdc++.h>
#define pb push_back
#define RAN(v)v.begin(),v.end()
#define FOR(i,v)\
	for(typeof(v.end())i=v.begin();i!=v.end();++i)
#define I (J+1)
#define J (i+j>>1)
#define P (k<<1)
#define S (P^1)
using namespace std;
const int inf=2e9;
struct buf{
	operator int(){
		int x=0,y=0,c=getchar();
		while(c<48)
			y=c==45,c=getchar();
		while(c>47)
			x=x*10+c-48,c=getchar();
		return y?-x:x;
	}
}it;
const int N=2e5+5;
struct info{int v,c;};
vector<info>e[N];
int ans=-inf;
int n,m,l,r,siz[N],w[N];
bool vis[N];
int len,f[N*4];
void pre(vector<int>&v,int i,int j,int k){
	if(i==j)f[k]=v[i-1];
	else{
		pre(v,i,J,P),pre(v,I,j,S);
		f[k]=max(f[P],f[S]);
	}
}
void pre(vector<int>&v){
	len=v.size();
	pre(v,1,len,1);
}
void cov(int u,int v,int i,int j,int k){
	if(i==j)f[k]=max(f[k],v);
	else{
		u<I?cov(u,v,i,J,P):cov(u,v,I,j,S);
		f[k]=max(f[P],f[S]);
	}
}
void cov(int u,int v){
	cov(u,v,1,len,1);
}
int ask(int u,int v,int i,int j,int k){
	return u==i&&j==v?f[k]:v<I?ask(u,v,i,J,P):u>J?ask(u,v,I,j,S):max(ask(u,J,i,J,P),ask(I,v,I,j,S));
}
int ask(int u,int v){
	u=max(u,1);
	v=min(v,len);
	return u>v?-inf:ask(u,v,1,len,1);
}
bool cmp(const vector<int>&a,const vector<int>&b){
	return a.size()>b.size();
}
void dfs1(int u,int p){
	siz[u]=1;
	FOR(i,e[u])
		if(i->v!=p&&!vis[i->v])
			dfs1(i->v,u),siz[u]+=siz[i->v];
}
int dfs2(int u,int p,int n){
	int s=n-siz[u];
	FOR(i,e[u])
		if(i->v!=p&&!vis[i->v]){
			if(int c=dfs2(i->v,u,n))return c;
			s=max(s,siz[i->v]);
		}
	return s*2<=n?u:0;
}
int dfs4(int u,int p,int d){
	int s=d;
	if(d<r)
		FOR(i,e[u])
			if(i->v!=p&&!vis[i->v])
				s=max(s,dfs4(i->v,u,d+1));
	return s;
}
vector<int>cur;
void dfs5(int u,int p,int c,int d,int s){
	cur[d]=max(cur[d],s);
	if(d<r)
		FOR(i,e[u])
			if(i->v!=p&&!vis[i->v])
				dfs5(i->v,u,i->c,d+1,c!=i->c?s+w[i->c]:s);
}
void dfs3(int u){
	dfs1(u,0);
	if(siz[u]>l){
		u=dfs2(u,0,siz[u]);
		vis[u]=1;
		map<int,vector<vector<int> > >col;
		FOR(i,e[u])
			if(!vis[i->v]){
				int n=dfs4(i->v,0,1);
				cur.assign(n,-inf);
				dfs5(i->v,0,i->c,0,w[i->c]);
				col[i->c].pb(cur);
			}
		vector<vector<int> >tmp;
		FOR(i,col){
			sort(RAN(i->second),cmp);
			vector<int>&s=i->second.front();
			pre(s);
			FOR(j,i->second)
				if(j!=i->second.begin()){
					vector<int>&v=*j;
					if(v.size()+len>=l)
						for(int k=0;k<v.size();++k){
							int a=ask(l-k-1,r-k-1);
							if(a!=-inf)
								ans=max(ans,a+v[k]-w[i->first]);
						}
					for(int k=0;k<v.size();++k){
						cov(k+1,v[k]);
						s[k]=max(s[k],v[k]);
					}
				}
			ans=max(ans,ask(l,r));
			tmp.pb(s);
		}
		if(tmp.size()>1){
			sort(RAN(tmp),cmp);
			pre(tmp.front());
			FOR(i,tmp)
				if(i!=tmp.begin()){
					vector<int>&v=*i;
					if(v.size()+len>=l)
						for(int j=0;j<v.size();++j){
							int a=ask(l-j-1,r-j-1);
							if(a!=-inf)
								ans=max(ans,a+v[j]);
						}
					for(int j=0;j<v.size();++j)
						cov(j+1,v[j]);
				}
		}
		FOR(i,e[u])
			if(!vis[i->v])dfs3(i->v);
	}
}
int main(){
	n=it,m=it,l=it,r=it;
	for(int i=1;i<=m;++i)
		w[i]=it;
	for(int i=2;i<=n;++i){
		int u=it,v=it,c=it;
		e[u].pb({v,c});
		e[v].pb({u,c});
	}
	dfs3(1);
	printf("%d\n",ans);
}