dp_energy 题解

前言

本题难度比较大，笔者现在还有些地方没弄懂，如果有不严谨的地方欢迎探讨纠正

题意

在一个 \(n\) 个节点的带权树中选择一个长度为 \(k\) 的连通块 \(S\)，使得 \(S\) 内节点之间的路径长度和加上非 \(S\) 内节点之间的路径长度和最大

形式化的答案：

\[\max{\frac{\sum_{i\in S} \sum_{j\in S} dis(i,j) + \sum_{i \notin S} \sum_{j\notin S} dis(i,j)}{2}} \]

题解

以下采用原题描述，且记 \(m = k\)

本题大部分人会想到 DP，并定义一个状态 \(f_{u,j}\) 表示以 \(u\) 为根的子树选 \(j\) 个黑点的答案，但这样明显无法转移

于是我们思考为什么不能转移，发现从 \(v\) 转移到 \(u\) 时，\(v\) 子树内节点到 \(u\) 会贡献一段距离，这个无法计算

那么我们能不能在转移 \(v\) 的时候就把这段距离算上？换而言之，我们可以计算 \(v\) 的子树时就把子树内边在全局的贡献加上。这样，我们从 \(v\) 转移到 \(u\) 时，就只用考虑边 \(v \rightarrow u\) 的贡献

但是还有一个连通块的限制，所以综上可得 DP 定义：\(f_{u,j,0/1}\) 表示 \(u\) 子树内选 \(j\) 个黑点，\(u\) 选或不选，子树内边在全局的贡献。

首先考虑 \(u\) 选的情况，根据树形背包，有转移方程

\[f_{u,j,1} = \max\{f_{u,j,1},f_{u,j-k,1}+f_{v,k,p(k)}+tot \times w_{v \rightarrow u}\} \]

其中

\[p(k) = \begin{cases} 0 & k=0 \\ 1 & k>0 \end{cases}\]

\(tot\) 为边被经过的次数，我们可以以此边为中心，把子树分为两边，则有 \(tot=\) 黑点贡献 \(+\) 白点贡献 \(= k \times (m-k) + (size_v-k) \times (n-m-(size_v-k))\)

有很多细节：

1. 应边计算 \(size\) 边转移，可以理解为每次做背包时加入一个子树
2. \(j,k\) 的上下界需注意，可以排除一些非法情况，并且把时间复杂度从 \(O(n^3)\) 优化到 \(O(n^2)\)，原因笔者尚未搞懂
3. \(f\) 数组应初始化为 \(- \infty\)，然后设初值，这样转移时可以排除非法情况

这部分代码如下：

F_(j,min(si[u],m),1){
	F(k,max(0,j-(si[u]-si[v])),min(si[v],j-1)){
		ll val = 1ll*k*(m-k) + 1ll*(si[v]-k)*(n-si[v]-m+k);
		dp[u][j][1] = max(dp[u][j][1], dp[u][j-k][1]+(k==0?dp[v][0][0]:dp[v][k][1])+val*edge[i].len);
	}
}

接下来是 \(u\) 不选的情况，由于 \(S\) 为连通块，所以 \(u\) 必只有一个儿子的子树有 \(j\) 个黑点，其他儿子的子树都没有黑点。这里，我们可以先算出儿子子树内没有黑色节点的贡献和，用 \(v\) 转移时加上总和，减去 \(v\) 的原贡献，再加上新贡献

最后给出代码

#include <bits/stdc++.h>
using namespace std;
#define F(i,a,b) for(int i = a; i <= b; i++)
#define F_(i,a,b) for(int i = a; i >= b; i--)
#define pii pair<int,int>
#define fr first
#define sc second
#define mem0(a) memset(a,0,sizeof(a))
#define pb push_back
#define lb(x) (x&(-x))
#define lson u<<1
#define rson u<<1|1
typedef long long ll;
typedef unsigned long long ull;
const int N = 2e3+10,M = 1e6;
const double eps = 1e-9;
const ll Mod = 1e9+7;
// huangyuze
int n,m;
struct EDGE{
	int from,to,pre;
	ll len;
} edge[N*2];
int final[N],ec,si[N];
ll dp[N][N][2];
void add(int u,int v,ll l){
	edge[++ec] = {u,v,final[u],l};
	final[u] = ec;
}
void dfs(int u,int fu){
	si[u] = 1;
	dp[u][1][1] = 0;
	ll sum = 0;
	for (int i=final[u]; i; i=edge[i].pre){
		int v = edge[i].to;
		if (v==fu) continue;
		dfs(v,u);
		sum += edge[i].len*si[v]*(n-m-si[v]) + dp[v][0][0];
	}
	dp[u][0][0] = sum;
	for (int i=final[u]; i; i=edge[i].pre){
		int v = edge[i].to;
		if (v==fu) continue;
		si[u] += si[v];
		
		F_(j,min(si[v],m),1){
			ll x = edge[i].len * (j*(m-j) + 1ll*(si[v]-j)*(n-si[v]-m+j));
			dp[u][j][0] = max(dp[u][j][0],max(dp[v][j][0],dp[v][j][1]) + sum - edge[i].len*si[v]*(n-m-si[v])-dp[v][0][0] + x);
		}
		
		F_(j,min(si[u],m),1){
			F(k,max(0,j-(si[u]-si[v])),min(si[v],j-1)){
				ll val = 1ll*k*(m-k) + 1ll*(si[v]-k)*(n-si[v]-m+k);
				dp[u][j][1] = max(dp[u][j][1], dp[u][j-k][1]+(k==0?dp[v][0][0]:dp[v][k][1])+val*edge[i].len);
			}
		}
	}
}
int main(){
	scanf("%d%d",&n,&m);
	F(i,1,n) F(j,0,m) dp[i][j][0] = dp[i][j][1] = -1e18;
	F(i,1,n-1){
		int u,v;
		ll w;
		scanf("%d%d%lld",&u,&v,&w);
		add(u,v,w),add(v,u,w);
	}
	dfs(1,0);
	printf("%lld",max(dp[1][m][1],dp[1][m][0]));
	return 0;
}