[BJOI2017] 树的难题

这道题目卡常卡了两个半小时仍然没有卡过。。。等进队了让队友帮忙卡一下吧

主要想一下思路

最主要的就是在计算路径长度的时候，假设当前递归到了点$i$，那么从点$i$出发的两条路径合并在一起，如果第一条边的颜色相同的话就会重复计算，为了解决这个问题，我们只用对每个点进行排序，将相同颜色的点放在一起，处理相同颜色的点内部，以及不同颜色的点之间的路径就好了，代码如下

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=2e5+10;
const ll INF=1e12;
int n,m,L,R;
ll c[N],maxdis[N];
ll Max[N<<2][2];
int siz[N],maxx[N];
struct node
{
	int num,color;
}temp;
vector<node> G[N];
int tag[2][N],New[N],TAG[2],NEW;
bool vis[N],ins[N],tf[2][N];
int sum,rt;
ll ans;
ll dist[N];
int len[N];
inline ll mymax(ll a,ll b)
{
	return a>b?a:b;
}
void calcsiz(int x,int fa) 
{
  	siz[x]=1;
  	maxx[x]=0;
  	for(register int j=0;j<G[x].size();j++)
    if(G[x][j].num!=fa&&!vis[G[x][j].num]) 
	{
      	calcsiz(G[x][j].num,x);
      	maxx[x]=mymax(maxx[x],siz[G[x][j].num]);
      	siz[x]+=siz[G[x][j].num];
    }
  	maxx[x]=mymax(maxx[x],sum-siz[x]);
  	if(maxx[x]<maxx[rt]) rt=x;
}
ll ask(int p,int l,int r,bool op,int x,int y)
{
	if(l>y||r<x) return -INF;
	if(l>=x&&r<=y) return Max[p][op];
	int mid=l+r>>1;
	return mymax(ask(p<<1,l,mid,op,x,y),ask(p<<1|1,mid+1,r,op,x,y));
}
void calcdist(int x,int fa,int col) 
{
	maxdis[len[x]]=mymax(maxdis[len[x]],dist[x]);
	//这里卡常，然而却没有什么用 
	if(!ins[len[x]])
	{
		New[NEW++]=len[x];
		ins[len[x]]=1;
		if(!tf[0][len[x]])
		{
			tf[0][len[x]]=1;
			tag[0][TAG[0]++]=len[x];
		}
	}
	if(len[x]<=R)
	ans=mymax(ans,dist[x]+ask(1,1,n+1,0,mymax(L-len[x],0)+1,R-len[x]+1));
  	for(register int j=0;j<G[x].size();j++)
    if(G[x][j].num!=fa&&!vis[G[x][j].num])
    {
    	len[G[x][j].num]=len[x]+1;
    	if(col!=G[x][j].color) dist[G[x][j].num]=dist[x]+c[G[x][j].color];
    	else dist[G[x][j].num]=dist[x];
		calcdist(G[x][j].num,x,G[x][j].color);
	} 
}
void modify(int p,int l,int r,bool op,int x,ll d)
{
	if(l>x||r<x) return;
	if(l==r)
	{
		if(d!=-INF) Max[p][op]=mymax(d,Max[p][op]);
		else Max[p][op]=d;
        //这里千万注意，如果是还原的话可以直接赋值，但是如果是更新的话一定要先比较
		return;
	}
	int mid=l+r>>1;
	modify(p<<1,l,mid,op,x,d);
	modify(p<<1|1,mid+1,r,op,x,d);
	Max[p][op]=mymax(Max[p<<1][op],Max[p<<1|1][op]);
}
bool cmp(node i,node j)
{
	if(vis[i.num]==vis[j.num]) return i.color<j.color;
	else return vis[i.num]<vis[j.num];
}
void dp(int x,int fa,int col)
{
	maxdis[len[x]]=mymax(maxdis[len[x]],dist[x]);
	if(!ins[len[x]])
	{
		New[NEW++]=len[x];
		ins[len[x]]=1;
	}
	if(len[x]<=R)
	ans=mymax(ans,dist[x]+ask(1,1,n+1,1,mymax(L-len[x],0)+1,R-len[x]+1)-c[col]);
  	for(register int j=0;j<G[x].size();j++)
    if(G[x][j].num!=fa&&!vis[G[x][j].num])
    dp(G[x][j].num,x,col);
}
void dfs(int x,int fa)
{
	tf[0][0]=1;
	tag[0][TAG[0]++]=0;
  	modify(1,1,n+1,0,1,0); 
  	vis[x]=1;
  	sort(G[x].begin(),G[x].end(),cmp);
  	//将颜色相同的放在一起 
  	for(register int j=0;j<G[x].size();)
    if(G[x][j].num!=fa&&!vis[G[x][j].num]) 
	{
      	int k=j;
      	while(j<G[x].size()&&!vis[G[x][j].num]&&G[x][k].color==G[x][j].color)
      	{
      		len[G[x][j].num]=1,dist[G[x][j].num]=c[G[x][j].color];
      		calcdist(G[x][j].num,x,G[x][j].color);
      		j++;
		}//找出同一颜色的儿子 
		for(register int o=0;o<NEW;o++) modify(1,1,n+1,0,New[o]+1,maxdis[New[o]]);
		for(register int o=0;o<NEW;o++)
		maxdis[New[o]]=-INF;
		for(register int o=0;o<NEW;o++)
		ins[New[o]]=0;
		NEW=0;
		for(register int o=k;o<j;o++)//处理同一颜色的点 
		{
			dp(G[x][o].num,x,G[x][o].color);
			for(register int w=0;w<NEW;w++) modify(1,1,n+1,1,New[w]+1,maxdis[New[w]]);
			for(register int w=0;w<NEW;w++)
			if(!tf[1][New[w]])
			{
				tf[1][New[w]]=1;
				tag[1][TAG[1]++]=New[w];
			}
			for(register int w=0;w<NEW;w++)
			maxdis[New[w]]=-INF;
			for(register int w=0;w<NEW;w++)
			ins[New[w]]=0;
			NEW=0;
		}
		for(register int o=0;o<TAG[1];o++)
		modify(1,1,n+1,1,tag[1][o]+1,-INF);
		for(register int o=0;o<TAG[1];o++)
		tf[1][tag[1][o]]=0;
		TAG[1]=0;
    }
    else break;
    for(register int j=0;j<TAG[0];j++)
	modify(1,1,n+1,0,tag[0][j]+1,-INF);
	for(register int j=0;j<TAG[0];j++)
	tf[0][tag[0][j]]=0;
	TAG[0]=0;
  	for(register int j=0;j<G[x].size();j++)
    if(G[x][j].num!=fa&&!vis[G[x][j].num]) 
	{
      sum=siz[G[x][j].num];
      rt=0;
      maxx[rt]=n+1;
      calcsiz(G[x][j].num,x);
      calcsiz(rt,0);
      dfs(rt,x);
    }
}
void build(int p,int l,int r)
{
	if(l==r)
	{
		for(register int i=0;i<=1;i++) 
		Max[p][i]=-INF;
		return;
	}
	int mid=l+r>>1;
	build(p<<1,l,mid);
	build(p<<1|1,mid+1,r);
	for(register int i=0;i<=1;i++) 
	Max[p][i]=-INF;
}
int read()
{
    int x=0,f=1;char s=getchar();
    while(s<'0'||s>'9'){if(s=='-')f=-f;s=getchar();}
    while(s>='0'&&s<='9'){x=x*10+s-48;s=getchar();}
    return x*f;
}
int main() 
{
  	n=read(),m=read(),L=read(),R=read();
  	for(register int i=1;i<=n;i++) maxdis[i]=-INF;
  	build(1,1,n+1);
  	//注意，线段树的下标表示长度，而且整体加了一 
  	//0表示不同颜色的线段树，1表示相同颜色的线段树 
  	for(register int i=1;i<=m;i++)
  	c[i]=read();
  	for(register int i=1,a,b,c;i<n;i++)
    {
    	a=read(),b=read(),c=read();
    	temp.num=b,temp.color=c;
		G[a].push_back(temp);
		temp.num=a;
		G[b].push_back(temp);
	}
  	rt=0;
  	maxx[rt]=n+1;
  	sum=n;
  	calcsiz(1,0);
  	calcsiz(rt,0);
  	ans=-INF;
  	dfs(rt,0);
  	printf("%lld",ans);
  	return 0;
}