洛谷 P3714 - [BJOI2017]树的难题（点分治）

洛谷题面传送门

咦？鸽子 tzc 竟然来补题解了？incredible（

首先看到这样类似于路径统计的问题我们可以非常自然地想到点分治。每次我们找出每个连通块的重心 \(x\) 然后以 \(x\) 为根 DFS 一遍整个子树，我们假设 \(y\) 到 \(x\) 的距离为 \(dep_y\)，\(x\to y\) 这一段上颜色的权值之和为 \(sum_y\)，那么考虑怎样合并两条路径。显然对于两个在 \(x\) 不同子树内的点 \(y,z\)，\(y\to z\) 路径上边的个数就是 \(dep_y+dep_z\)，路径上权值之和就是 \(sum_y+sum_z-c_{col_y}·[col_y=col_z]\)，其中 \(col_y\) 为 \(x\to y\) 路径上经过的第一条边的权值。看到这个 \([col_y=col_z]\) 貌似有点棘手，不过注意到我们贡献显然是一个子树一个子树计算的对吧，因此我们考虑将 \(x\) 所有子树按 \(x\) 到这棵子树经过的第一条边的颜色从小到大排序，然后维护两棵线段树，第一棵线段树上下标为 \(d\) 的位置上维护 \(\max\limits_{dep_y=d\land col_y\ne C}sum_y\)，第二棵维护 \(\max\limits_{dep_y=d\land col_y=C}sum_y\)，其中 \(C\) 为当前颜色种类，然后每次颜色改变就暴力地将第二棵线段树中所有元素插入第一棵线段树中即可，查询就在两棵树中分别查 \([r-dep_y,l-dep_y]\) 的最大值，记作 \(mx1\) 和 \(mx2\)，然后用 \(mx1+sum_y,mx2-c_{col_y}+sum_y\) 更新答案即可。

时间复杂度 \(n\log^2n\)，其中一个 \(\log\) 在于点分治，一个 \(\log\) 在于线段树。

最后稍微总结一下这类点分治解决树上路径计数题目的解题技巧：首先要考虑怎样合并两段路径，如果不好合并那一般使用点分治不太好解决，其次要思考如何维护两段路径的决策，比较简单的使用一个桶即可维护，比较复杂的需用 BIT/线段树/平衡树解决。对于计数问题，容斥也是一个不错的选择，即先不考虑两个端点不在同一子树这一条件，先一股脑把贡献全加上去，再扣掉在同一子树内的情况。

const int MAXN=2e5;
const int INF=0x3f3f3f3f;
const ll INFll=0x3f3f3f3f3f3f3f3fll;
int n,m,L,R,c[MAXN+5],hd[MAXN+5],to[MAXN*2+5],val[MAXN*2+5],nxt[MAXN*2+5],ec=0;
void adde(int u,int v,int w){to[++ec]=v;val[ec]=w;nxt[ec]=hd[u];hd[u]=ec;}
int mx[MAXN+5],cent=0,siz[MAXN+5];bool vis[MAXN+5];
void findcent(int x,int f,int tot){
	mx[x]=0;siz[x]=1;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f||vis[y]) continue;
		findcent(y,x,tot);siz[x]+=siz[y];
		chkmax(mx[x],siz[y]);
	} chkmax(mx[x],tot-siz[x]);
	if(mx[cent]>mx[x]) cent=x;
}
int dep[MAXN+5];ll sum[MAXN+5];
void getdep(int x,int f,int pre){
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e],z=val[e];if(y==f||vis[y]) continue;
		dep[y]=dep[x]+1;sum[y]=sum[x]+((z==pre)?0:c[z]);
		getdep(y,x,z);
	}
}
struct segtree{
	struct node{int l,r;ll mx;} s[MAXN*4+5];
	stack<int> stk;
	void pushup(int k){s[k].mx=max(s[k<<1].mx,s[k<<1|1].mx);stk.push(k);}
	void build(int k,int l,int r){
		s[k].l=l;s[k].r=r;s[k].mx=-INFll;if(l==r) return;
		int mid=l+r>>1;build(k<<1,l,mid);build(k<<1|1,mid+1,r);
	}
	void modify(int k,int p,ll v){
		if(s[k].l==s[k].r) return chkmax(s[k].mx,v),stk.push(k),void();
		int mid=s[k].l+s[k].r>>1;(p<=mid)?modify(k<<1,p,v):modify(k<<1|1,p,v);
		pushup(k);
	}
	ll query(int k,int l,int r){
		if(l>r) return -INFll;
		if(l<=s[k].l&&s[k].r<=r) return s[k].mx;
		int mid=s[k].l+s[k].r>>1;
		if(r<=mid) return query(k<<1,l,r);
		else if(l>mid) return query(k<<1|1,l,r);
		else return max(query(k<<1,l,mid),query(k<<1|1,mid+1,r));
	}
	void relax(){
		while(!stk.empty()){
			int k=stk.top();stk.pop();
			s[k].mx=-INFll;
		}
	}
} s1,s2;
vector<int> pt;
ll res=-INFll;
void findpts(int x,int f){
	pt.pb(x);
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(y==f||vis[y]) continue;
		findpts(y,x);
	}
}
void divcent(int x){
//	printf("divcent %d\n",x);
	vis[x]=1;dep[x]=sum[x]=0;
	vector<pii> sub;
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e],z=val[e];if(vis[y]) continue;
		dep[y]=1;sum[y]=c[z];getdep(y,x,z);
		sub.pb(mp(z,y));
	} sort(sub.begin(),sub.end());
	vector<int> wt;s1.modify(1,0,0);
	for(int i=0;i<sub.size();i++){
		if(i&&sub[i].fi!=sub[i-1].fi){
			s2.relax();
			for(int y:wt) s1.modify(1,dep[y],sum[y]);
			wt.clear();
		} int y=sub[i].se;pt.clear();findpts(y,x);
		for(int z:pt){
			int d=dep[z];
			if(d<=R){
				chkmax(res,s1.query(1,max(L-d,0),R-d)+sum[z]);
				chkmax(res,s2.query(1,max(L-d,0),R-d)+sum[z]-c[sub[i].fi]);
			}
		} for(int z:pt) wt.pb(z),s2.modify(1,dep[z],sum[z]);
//		for(int z:pt) printf("%d %d %lld\n",z,dep[z],sum[z]);
	} s1.relax();s2.relax();
	for(int e=hd[x];e;e=nxt[e]){
		int y=to[e];if(vis[y]) continue;cent=0;
		findcent(y,x,siz[y]);divcent(cent);
	}
}
int main(){
	scanf("%d%d%d%d",&n,&m,&L,&R);s1.build(1,0,n);s2.build(1,0,n);
	for(int i=1;i<=m;i++) scanf("%d",&c[i]);
	for(int i=1,u,v,w;i<n;i++) scanf("%d%d%d",&u,&v,&w),adde(u,v,w),adde(v,u,w);
	mx[0]=INF;findcent(1,0,n);divcent(cent);
	printf("%lld\n",res);
	return 0;
}