【贪心】AcWing 115. 给树染色

传送门：https://www.acwing.com/solution/content/82014/

分析

假如没有染色顺序的约束，那么最佳决策当然是先染权值大的点（本质上就是排序不等式）。

然而现在它有约束，但我们可以保证的一点是：当树上最大的点 \(u\) 的父节点 \(p\) 被染色的时候，立刻将 \(u\) 进行染色是最优的：这意味着，我们可以将其合并成一个新的点，等价于染色步骤中连续对这两个点进行染色。

那么这个新点的权值应该是多少呢？考虑现在 \(X,Y\) 点组成一个合并的点（对应上述的 \(p, u\)）\(X\) 可以被染色，然后 \(Z\) 点也是当前可以被染色的点。（记它们点权为自己的小写字母）

假设现在已经染了 \(k\) 个点

• 如果先染 \((X, Y)\)，那么贡献是 \((k+1)x + (k+2)y + (k+3)z\)。
• 如果先染 \(Z\)，贡献为 \((k+1)z + (k+2)x + (k+3)u\)

也就是合并的 \((X, Y)\) 和 \(Z\) 染色的优先级取决于两个贡献的大小。

\((k+1)x + (k+2)y + (k+3)z < (k+1)z + (k+2)x + (k+3)y\) 等价于 \((x + y)/2 > z\)，即 \((x + y)/2 > z\) 时选择先染 \(Z\) 否则先染 \((X,Y)\)。

据此，将一个新点权值设置为两点平均值是不改变决策的优先级的，推而广之，如果一个新点是由两个新点合并而来，那么权值就是两点的权值和的平均值（例如一个新点 \(A\) 有 \(a\) 个起初树上的点，\(B\) 有 \(b\) 个起初树上的点，那么权值就是 \((sum_A + sum_B) / (a + b)\)

实现

y总 的做法是利用偏移量进行统计，感觉不太好想 qwq，所以我采取的是将合并的过程进行模拟，并记录相关信息，最后对剩下那个点的信息进行展开得到决策序列（有点像解压缩）。

写这题的时候心很慌，感觉自己的实现挺乱的（感觉就我是这样乱搞的），没想到能 1A。如果是因为数据水了可以来试试 hack qwq。

看代码：

// Problem: 给树染色
// Contest: AcWing
// URL: https://www.acwing.com/problem/content/117/
// Memory Limit: 64 MB
// Time Limit: 1000 ms
// 
// Powered by CP Editor (https://cpeditor.org)

#include<bits/stdc++.h>
using namespace std;

#define debug(x) cerr << #x << ": " << (x) << endl
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define dwn(i,a,b) for(int i=(a);i>=(b);i--)
#define pb push_back
#define all(x) (x).begin(), (x).end()

using pii = pair<int, int>;
using ll = long long;

inline void read(int &x){
    int s=0; x=1;
    char ch=getchar();
    while(ch<'0' || ch>'9') {if(ch=='-')x=-1;ch=getchar();}
    while(ch>='0' && ch<='9') s=(s<<3)+(s<<1)+ch-'0',ch=getchar();
    x*=s;
}

const int N=2022;

int n, rt;
int w[N];

struct Node{
	int fa, sz, sum;
	double avg;
}e[N];

int tot;
int h[N], ne[N]; // 记 cur=(p, u), 那么 h[cur]=p, ne[cur]=u
bool vis[N], inr[N]; // 分别代表：是否已经被并入新点，是否在根节点所在的新点中

int find(){
	int u; double val=0;
	rep(i,1,tot) if(!vis[i] && !inr[i] && e[i].avg>val) val=e[i].avg, u=i;
	return u;
}

vector<int> get(int cur){ // 解压缩
	int p=h[cur], u=ne[cur];
	vector<int> v1, v2;
	if(p>n) v1=get(p);
	else v1={p};
	
	if(u>n) v2=get(u);
	else v2={u};
	
	v1.insert(end(v1), all(v2));
	return v1;
}

vector<int> g[N];

void getFa(int u, int fa){
	e[u].fa=fa;
	for(auto go: g[u]) if(go!=fa) getFa(go, u);
}

int main(){
	cin>>n>>rt;
	rep(i,1,n) read(w[i]), e[i].sum=e[i].avg=w[i], e[i].sz=1;
	
	inr[rt]=1;
	rep(i,1,n-1){
		int u, v; read(u), read(v);
		g[u].pb(v), g[v].pb(u);
	}
	
	getFa(rt, 0); // 将每个点的父亲处理出来
	
	tot=n;
	rep(_,1,n-1){
		int u=find(), p=e[u].fa; // 找到 u 以及 p
		vis[u]=vis[p]=1; // 标记为已合并

		tot++;
		e[tot]={e[p].fa, e[u].sz+e[p].sz, e[u].sum+e[p].sum, (double)(e[u].sum+e[p].sum)/(e[u].sz+e[p].sz)};
		h[tot]=p, ne[tot]=u;
		if(inr[p]) inr[u]=inr[tot]=1;
		rep(i,1,tot-1) if(e[i].fa==u || e[i].fa==p) e[i].fa=tot; // 维护发生修改的每个点父节点信息
	}

	auto vec=get(tot); // 从合并得到的最后的新点（此时树上只剩一个点）开始解压缩得到合并的操作序列。
	
	int res=0;
	rep(i,1,n) res+=i*w[vec[i-1]];
	cout<<res<<endl;
	
	return 0;
}