AcWing 115. 给树染色

AcWing 115. 给树染色(https://www.acwing.com/problem/content/description/117/)
一颗树有 n 个节点，这些节点被标号为：1,2,3…n，每个节点 i 都有一个权值 A[i]。
现在要把这棵树的节点全部染色，染色的规则是：
根节点 R 可以随时被染色；对于其他节点，在被染色之前它的父亲节点必须已经染上了色。

每次染色的代价为 T×A[i]，其中 T 代表当前是第几次染色。

求把这棵树染色的最小总代价。

输入格式
第一行包含两个整数 n 和 R，分别代表树的节点数以及根节点的序号。

第二行包含 n 个整数，代表所有节点的权值，第 i 个数即为第 i 个节点的权值 A[i]。

接下来 n−1 行，每行包含两个整数 a 和 b，代表两个节点的序号，两节点满足关系： a 节点是 b 节点的父节点。

除根节点外的其他 n−1 个节点的父节点和它们本身会在这 n−1 行中表示出来。

同一行内的数用空格隔开。

输出格式
输出一个整数，代表把这棵树染色的最小总代价。

数据范围
1≤n≤1000,
1≤A[i]≤1000
输入样例：
5 1
1 2 1 2 4
1 2
1 3
2 4
3 5
输出样例：
33

贪心，嗯，怎么个贪法呢？显然不能从根开始每步选最大的权值染色，因为后面的节点权值大小未知。但整体上来看，显然我们要先找到权值最大的节点，但是想要给这个节点染色，我们还必须得给他的根节点染色，层层向上，直到根节点。我们选择的策略应该是什么呢？
假定有x,y,z三个节点，且我们已知x,y节点有父子关系，染色过程是连续的，则我们可以选择先染色x,y再染z，这种情况总花费是x+2y+3z,另一种是先染z再染x,y，这种情况总花费是z+2x+3y,假设第一种染色方案更优，则有：x+2y+3z-(z+2x+3y)<0，即：z<(x+y)/2,这给了我们提示，想要总花费最小，我们需要优先将某几个平均权值最大的节点连续染色，故我们可以从非根节点开始，每次找到平均权值最大的节点(初始时所有节点平均权值都是原来的权值)，将其和父节点合并成一个节点并更新父节点的相关参数，最终整棵树会合并到根节点上，并且整个合并过程可以直到合并顺序，
每次寻找最大平均权值时间复杂度是On，整体On2，如果维护一个二叉堆，时间复杂度可以降为Onlgn。
PS:在我苦想半天怎么记录节点合并顺序无果后看了y老师的代码以及思路，只能说y老师太强了，在这里附上y老师的题解：https://www.acwing.com/solution/content/1065/，y老师并没有将整个合并顺序先记下来，而是通过合并的过程中不断将节点权值*父节点上合并节点个数累加到结果中，父节点中已合并节点的个数就指示了在染当前节点前我们已经给多少个节点染色。
AC代码：

#include<bits/stdc++.h>
using namespace std;
struct nd
{
	double avg;
	int sum=0;
	int fa;
	int val;
}a[1005];
int main()
{
	int n,r;
	int x,y;
	while (cin>>n>>r&&n)    //此步忽略，学校OJ上的数据和AcWing上不太一样，不过在哪交都是对的，读到EOF自然会结束循环
	{
		int res=0;
		for(int i=1;i<=n;i++) {
			cin>>a[i].val;
			a[i].avg=a[i].val;
			a[i].sum=1;
			res+=a[i].val;
		}
		for(int i=1;i<n;i++)
		{
			cin>>x>>y;
			a[y].fa=x;
		}
		for(int i=1;i<n;i++){	//n-1次合并
			double mavg=-1;
			int key;
			//找平均权值最大节点
			for(int j=1;j<=n;j++){
				if(j!=r&&a[j].avg>mavg){
					mavg=a[j].avg;
					key=j;
				}
			}
			int father=a[key].fa;
			res+=a[key].val*a[father].sum;
			a[father].sum+=a[key].sum;
			a[father].val+=a[key].val;
			a[father].avg=a[father].val*1.0/a[father].sum;
			//将该节点所有子节点的父节点更新
			for(int j=1;j<=n;j++){
				if(a[j].fa==key)
					a[j].fa=father;
			}
			a[key].avg=-0x3f3f;		//合并结束节点权值改变避免被重复遍历到
		}
		cout<<res<<endl;
	}
	
    return 0; 
}

---

果然，菜鸟的我还是只学会看完别人代码后大呼：妙啊~