Paint Tree 2

\(Paint\ Tree\ 2\)

F - Paint Tree 2

题意：

给定一棵\(n\)个结点的树，点权为\(a_i\)，选取至多\(k\)条不相交的链，计算链点权和的最大值。

\(1\le n\le2\times 10^5,1\le k\le 5\)

题解：

发现\(k\)很小，考虑用类似树上背包的\(dp\)解决。首先考虑状态\(dp_{u,k}\)为\(u\)为根的子树中选取了\(k\)条链的最大值，但是这样无法转移，因为我们无法确定链的形态。

而对于点\(u\)的形态一般考虑，

\(0.\)\(u\)不在任意链中。

\(1.\)\(u\)是链的端点。

\(2.\)\(u\)不是链的端点而是内部的点。

状态\(s\)的定义同上。

将状态修改为\(dp_{u,k,s}\)，转移时一般考虑，

\(1.\)不连接\((u,v)\)这条边，直接继承\(v\)相关状态。

\(2.\)连接\((u,v)\)这条边，且让\(v\)的链延伸到\(u\)。

\(3.\)连接\((u,v)\)这条边，且让\(v\)的链与\(u\)的链在\(u\)处合并。

综上，一般考虑顺序为，\((u,v)\)是否连接？\(u,v\)链是延伸还是合并？

转移方程如下：

不连接\((u,v)\)时\(u\)直接继承\(v\)的状态：

\(dp_{u,i,s}=max\{dp_{u,i,s},dp_{u,i-j,s}+max\{dp_{v,j,p},p=0,1,2\}\}\)

连接\((u.v)\)且让\(v\)的链延伸到\(u\)：

\(dp_{u,i,1}=max\{dp_{u,i,1},dp_{u,i-j,0}+dp_{v,j,1}+a_u \}\)

连接\((u,v)\)且让\(u\)的链和\(v\)的链在\(u\)合并，即两条链合并成一条更长的链，链的数量少了\(1\)，因此\(u\)有\(i-j\)条链，\(v\)有\(j+1\)条合并成\(i\)条链：

\(dp_{u,i,2}=max\{dp_{u,i,2},dp_{u,i-j,1}+dp_{v,j+1,1} \}\)

初值为\(dp_{u,0,0}=0，dp_{u,1,1}=dp_{u,1,2}=a_u\)。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
const ll inf=1e18;
const int maxn=2e5+5;
int n,k;
ll a[maxn],dp[maxn][6][3];
vector<int>g[maxn];
void dfs(int u,int fa){
	dp[u][0][0]=0;
	dp[u][1][1]=dp[u][1][2]=a[u];
	for(auto v:g[u]){
		if(v==fa) continue;
		dfs(v,u);
		for(int i=k;i>=0;i--){
			for(int j=0;j<=i;j++){
				dp[u][i][0]=max(dp[u][i][0],dp[u][i-j][0]+max(dp[v][j][0],max(dp[v][j][1],dp[v][j][2])));
				dp[u][i][1]=max(dp[u][i][1],dp[u][i-j][1]+max(dp[v][j][0],max(dp[v][j][1],dp[v][j][2])));
				dp[u][i][2]=max(dp[u][i][2],dp[u][i-j][2]+max(dp[v][j][0],max(dp[v][j][1],dp[v][j][2])));
				dp[u][i][1]=max(dp[u][i][1],dp[u][i-j][0]+dp[v][j][1]+a[u]);
				dp[u][i][2]=max(dp[u][i][2],dp[u][i-j][1]+dp[v][j+1][1]);
			}
		}
	}
}
int main(){
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	cin>>n>>k;
	for(int i=1;i<=n;i++) cin>>a[i];
	for(int i=1;i<n;i++){
		int u,v;
		cin>>u>>v;
		g[u].push_back(v);
		g[v].push_back(u);
	}
	for(int i=1;i<=n;i++){
		for(int j=0;j<=k;j++){
			dp[i][j][0]=dp[i][j][1]=dp[i][j][2]=-inf;
		}
	} 
	dfs(1,0);
	ll ans=-inf;
	for(int i=0;i<=k;i++) ans=max(ans,max(dp[1][i][0],max(dp[1][i][1],dp[1][i][2])));
	cout<<ans<<endl;
	return 0;
}