AT_abc259_f [ABC259F] Select Edges 题解

分析

考虑树形 DP。

我们可以以节点 $1$ 为根，从叶子节点往上推。注意到，任意一个除根节点外的节点，其与其父节点都必有一条连边。所以在定义 DP 数组时，我们需要将其分成两类：

1. $f_{i,1}$：不选择 $i$ 与其父节点的连边时，以 $i$ 为根所得子树的最大价值。

2. $f_{i,0}$：选择 $i$ 与其父节点的连边时，以 $i$ 为根所得子树的最大价值。

这样，我们在往上推的时候，$f_{i,0}$ 与 $f_{i,1}$ 的区别就在与取不取 $d_i$ 条连向子节点的边。而 $f_{i,0}$ 与 $f_{i,1}$ 的最初价值，一定是在 $i$ 的任意子节点 $j$，满足：$f_{j,1}>0$ 时的总和。如果有一个子节点 $j$，它有：$f_{j,0}+w>f_{j,1}$（$w$ 是 $i$ 与 $j$ 的连边边权），那么我们一定是把原来的 $f_{j,1}$ 换成 $f_{j,0}+w$。这个增加量是可以求的，即：$f_{j,0}+w-f_{j,1}$。我们求出所有的增加量，从大到小排序后，使用前 $\min (d_i,k)$（$k$ 为增加量的数量）个增加量，当做修改值对 $f_{i,0}$ 与 $f_{i,1}$ 进行调整即可。

最后的答案就是：$\max (f_{1,0},f_{1,1})$。

注：为什么我们的最初价值可以是上面的那个和呢？因为我们的最初价值就相当于不取节点 $i$ 的任意连边的最大价值，而 $f_{j,1}$ 正好就不能取 $i$ 与 $j$ 的连边。又因为得到的最大价值可能为负数，所以我们还要看一下是不是利于求最大值的。

代码

#include<bits/stdc++.h>
using namespace std;
#define int long long
const int N=1e6+10;
int n,x[N];
int d[N],w[N],ne[N],e[N],h[N],idx;
void add(int a,int b,int c){
	ne[idx]=h[a],w[idx]=c,e[idx]=b,h[a]=idx++;
}
int f[N][2];
void dfs(int now,int fa){
	int now_ans=0,ans=0;
	priority_queue<int> qu;//大根堆，取价值大的
	for(int i=h[now];~i;i=ne[i]){
		int j=e[i];
		if(j==fa) continue;
		dfs(j,now);
		if(f[j][1]>0) now_ans+=f[j][1];
		if(d[j]){//子节点能连边
			if(f[j][0]+w[i]>f[j][1]){//更优
				qu.push(f[j][0]+w[i]-f[j][1]);//增加量
			}
		}
	}
	f[now][0]=f[now][1]=now_ans;//最初价值
	for(int i=1;i<=d[now];i++){
		if(!qu.empty()){
			ans+=qu.top();qu.pop();
		}
		if(i==d[now]-1){//留一条边给now->fa
			f[now][0]+=ans;
		}
		else if(i==d[now]){//不留
			f[now][1]+=ans;
		}
	}
}
int ans;
signed main(){
	memset(h,-1,sizeof(h));
	cin>>n;
	for(int i=1;i<=n;i++){
		cin>>d[i];
	}
	for(int i=1;i<n;i++){
		int a,b,c;cin>>a>>b>>c;
		add(a,b,c),add(b,a,c);
	}
	dfs(1,-1);
	ans=max(f[1][0],f[1][1]);
	return cout<<ans,0;
}