较难的换根dp:P6213 「SWTR-04」Collecting Coins

传送门

前题提要:感觉这道换根dp可以说是集中了换根dp的所有较高难度的操作和思想,以及较高的一些实现细节,可以说能够完全写出这道题才叫真正理解了换根dp,非常值得一做.

首先读完题意,不难发现这道题有很多限制.点的访问次数限制,必须访问某一个点,想要获得最大的贡献,没有限制初始起点.乍一看感觉完全没办法解决,我们考虑一个一个的去解决上述限制.

首先假设只想获得最大的贡献,那么这个显然是一个简单的树形dp,也就是$f[u]=max(f[v]+w)$.$f[u]$表示访问$u$子树所能获得的最大贡献.
然后我们考虑加上访问次数的限制,我们发现对于一个点,无论他是起始点还是中途经过的点,只要访问这个点,我们就损失了一个进入机会.然后显然我们想访问一个子树获得子树的贡献就需要花费一次进入机会,因为我们进入一个子树必然需要返回.诶,此时我们就可以解决了.假设一个点的访问次数限制为$k$,那么对于这个点来说,我们最多只能加入$k-1$个子树的贡献.所以我们考虑对所有子树的贡献进行一个排序,然后取前$k-1$个最大的贡献即可.
然后我们再考虑加入必须访问一个点的贡献.这一点直接解决起来就有点困难了,我们需要转化一下这个条件.我们考虑将这个必须访问的点当做我们的遍历的树的根节点,这样我们的必须访问的需求就迎刃而解了.因为对于刚开始的一棵树来说,我们这个根节点是必须遍历的,没有任何问题.
然后因为本题并没有限制初始点,所以我们需要进行一个换根.那么对于换完根的树来说,因为存在一个必须访问的点,所以我们限制换完根的子树必须选祖先那个子树即可.然后在自己的子树里选出前$k-2$大的.

---

至此本题的主要解决思路就结束了,但是本题在换根上的细节和经验问题还是比较多的(反正我当时在写这道题被细节方面搞了很久),接下来详细讲讲细节方面的处理

其实建议不看下文自己解决细节问题,毕竟只有自己的连续思考才毕竟清楚

考虑我们维护出了$f[u]$表示访问一个节点所能获得的最大贡献,并且因为换根需求,我们考虑记录使用$max\_f$记录下所有的u子树对u的贡献,并且对此从大到小排个序.
PS:这种记录子树对u的贡献的方法在很多题目中都需要体现(因为换根的时候需要考虑自身节点对父亲的贡献),但是很多题目都是记录最大值和次大值即可,这样较为方便,但是实际上那些题目都是可以直接记录所有的贡献然后等会处理,并且在本题中光记录最大和次大并不能解决问题,所以考虑直接记录所有的贡献

因为我们的父亲节点是必选的,为了方便解决问题,不妨设置一个辅助数组$fa\_val[u]$来记录以$u$为根时$u$节点的父亲子树对$u$的贡献.然后设$dp[u]$记录以$u$为根时的贡献.
先来维护$fa\_val[u]$数组
然后我们考虑根节点从$u->v$转移,此时需要进行分类讨论,考虑先维护辅助数组:

1. 当v子树是$f[u]$数组的一个贡献来源时:此时我们想要维护$fa\_val[u]$时,我们需要花费一次进入机会从$v$进入到$u$,然后我们还需要一次进入机会从$u$到$u$的父亲节点(因为父亲节点是必选的),所以此时我们就只剩下$k-2$棵可选择的子树贡献(并且这个k-2可子树是不能包括v的),那么考虑预处理出原本$u$子树原$k-1$课子树的总贡献,然后减掉$v$这个子树的贡献就维护出了前$k-2$棵除v子树外的最大贡献了.但是光光这样的方式是有问题的(当时卡了我很长时间),因为假设我们的u是根节点时,我们此时没有父亲那边的贡献,并且我们也不需要花一次进入机会去进入到父亲节点,所以这种情况下是需要特判的,具体特判方式看下文代码即可,此处不在赘述
2. 考虑v子树不是$f[u]$的一个贡献来源,那么此时我们想要维护$fa\_val[u]$时,我们同样此时只需要记录$u$子树除$v$外的前$k-2$棵子树的贡献即可.因为此时v不处于前$k-1$大的贡献之间,所以我们直接预处理前$k-2$的$max\_f$的和即可.同样的,当我们是处于根节点时照样需要特判.

考虑我们如果维护出了$fa\_val$数组之后,对于我们的$dp$数组就不难维护了,我们此时的$dp[u]=fa\_val[u]+u前k-2大的子树的贡献和$

PS:此处同时也有一个小trick,因为时间复杂度的约束,我们不能实时的枚举v是否是u的一个贡献子树,所以我们需要在遍历子树之前先预处理一下这个状态,并且同时记录下前k-1大,前k大,前k-2大的子树的贡献和即可

---

下面是具体的代码部分:

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define root 1,n,1
#define ls rt<<1
#define rs rt<<1|1
#define lson l,mid,rt<<1
#define rson mid+1,r,rt<<1|1
inline ll read() {
	ll x=0,w=1;char ch=getchar();
	for(;ch>'9'||ch<'0';ch=getchar()) if(ch=='-') w=-1;
	for(;ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
	return x*w;
}
inline void print(__int128 x){
	if(x<0) {putchar('-');x=-x;}
	if(x>9) print(x/10);
	putchar(x%10+'0');
}
#define maxn 1000000
const double eps=1e-8;
#define	int_INF 0x3f3f3f3f
#define ll_INF 0x3f3f3f3f3f3f3f3f
struct Edge{
	int v,w;
};
vector<Edge>edge[maxn];int k[maxn];
int f[maxn];vector<pair<int,int> >max_f[maxn];
int n,d;
void dfs1(int u,int per_u) {
	for(auto [v,w]:edge[u]) {
		if(v==per_u) continue;
		dfs1(v,u);
		max_f[u].push_back({f[v]+w,v});
	}
	sort(max_f[u].begin(),max_f[u].end());
	reverse(max_f[u].begin(),max_f[u].end());
	for(int i=0;i<min((int)max_f[u].size(),k[u]-1);i++) {
		f[u]+=max_f[u][i].first;
	}	
}
int dp[maxn];int check[maxn];int fa_val[maxn];
void dfs2(int u,int per_u) {
	int sum1=0,sum2=0,sum3=0;
	for(int i=0;i<max_f[u].size();i++) {
		int v=max_f[u][i].second;int w=max_f[u][i].first;
		if(i<k[u]-1) sum1+=w,check[v]=1;
		if(i<k[u]) sum2+=w;
		if(i<k[u]-2) sum3+=w;
	}
	dp[u]+=sum3;
	for(auto [v,w]:edge[u]) {
		if(v==per_u) continue;
		if(k[v]==1) continue;
		if(check[v]) {
			if(u!=d) 
				fa_val[v]=fa_val[u]+sum1-f[v]-w+w;
			else 
				fa_val[v]=fa_val[u]+sum2-f[v]-w+w;				
			dp[v]=fa_val[v];
		}
		else {
			if(u!=d) 
				fa_val[v]=fa_val[u]+sum3+w;
			else 
				fa_val[v]=fa_val[u]+sum1+w;	
			dp[v]=fa_val[v];
		}	
		dfs2(v,u);
	}
}
int main() {
	n=read();d=read();
	for(int i=1;i<n;i++) {
		int u=read();int v=read();int w=read();
		edge[u].push_back({v,w});
		edge[v].push_back({u,w});
	}
	for(int i=1;i<=n;i++) {
		k[i]=read();
	}
	dfs1(d,0);dfs2(d,0);
	int ans=f[d];
	for(int i=1;i<=n;i++) {
		ans=max(ans,dp[i]);
	}
	cout<<ans<<endl;
	return 0;	
}