题解 | #A.Tree# 2023牛客暑期多校6
A.Tree
图论-Kruskal、动态规划
题目大意
给定一棵 n n n 个节点带点权和边权的无根树
节点具有颜色白（ 0 0 0）和黑（ 1 1 1），颜色可反转，所需代价 c o s t i cost_i costi 为该点点权
整棵树的 e a r n i n g earning earning 为 ∑ u ∈ V 0 ∑ v ∈ V 1 v a l ( u , v ) \sum\limits_{u\in V_0}\sum\limits_{v\in V_1} val(u,v) u∈V0∑v∈V1∑val(u,v) 。其中， v a l ( u , v ) val(u,v) val(u,v) 为节点 u → v u\rightarrow v u→v的最短路径上的最大边权， V 0 V_0 V0 为白色点集， V 2 V_2 V2 为黑色点集
可以操作反转节点颜色任意次，求最大 e a r n i n g − ∑ c o s t earning-\sum cost earning−∑cost （以下简称 s c o r e score score ）
解题思路
注意到对于每对黑白点对，其贡献为最短路径上的最大边权。考虑利用Kruskal算法对树进行重构，即按边权从小到大的顺序进行加边
可以发现，在加入某条边 e i e_i ei 时，原本在同一连通分量中的黑白点对的 s c o r e score score 不受影响。由于加边顺序，新加的边一定具有目前最大的边权，因此只有经过新加的这条边的点对才对 s c o r e score score 具有贡献，每个点对的贡献值为 w i w_i wi ，点对数量为： 左白 × 右黑 + 左黑 × 右白 左白\times右黑+左黑\times右白 左白×右黑+左黑×右白
构造dp数组：定义 d p i , j dp_{i,j} dpi,j 为连通分量 i i i （以并查集中连通分量的根节点标识）中具有 j j j 个白色节点时的最大 s c o r e score score 。初始对于点 i i i ， d p i , c o l o r i = 0 dp_{i,color_i}=0 dpi,colori=0 （不变）； d p i , c o l o r i ⊕ 1 = − c o s t i dp_{i,color_i\oplus 1}=-cost_i dpi,colori⊕1=−costi （反转）
在合并 A 、 B A、B A、B 两个连通分量到 C C C 时具有以下转移方程：
d p C , i = max 0 ≤ k ≤ ∣ A ∣ & 0 ≤ i − k ≤ ∣ B ∣ { d p A , k + d p B , i − k + w ( k ( ∣ B ∣ − ( i − k ) ) + ( i − k ) ( ∣ A ∣ − k ) ) } dp_{C,i}=\max\limits_{0\le k \le |A| \And 0\le i-k\le |B|}\{dp_{A,k}+dp_{B,i-k}+w(k(|B|-(i-k))+(i-k)(|A|-k))\} dpC,i=0≤k≤∣A∣&0≤i−k≤∣B∣max{dpA,k+dpB,i−k+w(k(∣B∣−(i−k))+(i−k)(∣A∣−k))}
（两边原有的 s c o r e score score 加上过新加边的 s c o r e score score ）
可以结合代码注释进行理解
时间复杂度
O ( n 2 log n ) O(n^2\log n) O(n2logn)
参考代码
参考代码为已AC代码主干，其中部分功能需读者自行实现
struct edge{
int u,v;ll w;
edge():u(0),v(0),w(0){}
edge(int _u,int _v,ll _w): u(_u),v(_v),w(_w) {}
bool operator<(const edge &b)const{
return w dp[N];//dp_i,j:集合i有j个白色点时的score
vector Edge;
//vector> Edge;//>;
vector cost(N);
ll find_dsu(ll x){
return dsu[x]==x?x:dsu[x]=find_dsu(dsu[x]);
}//并查集
void merge_dsu(ll a,ll b,ll w){
a=find_dsu(a);b=find_dsu(b);
if(sz[a]>sz[b]) swap(a,b);
if(a==b) return ;
ll n=sz[a]+sz[b],rt;
//n:a,b点集体积之和，即白点数量上限
vector tmp(n+1);
FORLL(i,0,n){
rt=-INF;
for(ll k=0;k<=sz[a];k++) if(k<=i&&i-k<=sz[b]){//k:a中白点数量
ll cnt0=i-k,cnt1=sz[b]-(i-k);//cnt0:b中白点数量 cnt1:b中黑点数量
rt=max(rt,dp[a][k]+dp[b][i-k]+w*(k*cnt1+cnt0*(sz[a]-k)));
//dp=max(两点集已经具有的最大earning-cost+经过这条边的点对数*边权)
}
tmp[i]=rt;
}
dp[b]=tmp;
sz[b]+=sz[a];
dsu[a]=b;
}
void solve()
{
ll n,t;cin >> n;
Edge.resize(n-1);
ll u,v,w;
FORLL(i,1,n){
dp[i].resize(2);//接下来有需求会扩容
dsu[i]=i;sz[i]=1;
}
FORLL(i,1,n) cin >> color[i];
FORLL(i,1,n){
cin >> cost[i];
dp[i][color[i]]=(ll)0;
dp[i][color[i]^1]=(ll)-cost[i];
}
for(auto &e:Edge){
cin >> e.u >> e.v >> e.w;
//cin >> u >> v >> w;
//e.second.first=u;e.second.second=v;e.first=w;
}//存边
SORT(Edge);//按边权从小到大重构树
edge e;
//pair e;
FORLL(i,0,n-2){
e=Edge[i];
u=e.u;v=e.v;w=e.w;
//u=e.second.first;v=e.second.second;w=e.first;
merge_dsu(u,v,w);
}auto tmp=dp[find_dsu(1)];
cout << *max_element(ALL(tmp)) << endl;
}