AcWing 115. 给树染色(Java)
原题链接
困难
作者:
上杉
,
2021-04-30 13:10:29
,
所有人可见
,
阅读 515
import java.util.*;
public class Main{
static int n;
static int root;
static Node[] nodes;
static long res;
public static void main(String[] args){
Scanner input = new Scanner(System.in);
n = input.nextInt(); root = input.nextInt();nodes = new Node[n+1];
for(int i = 1 ; i <= n ;i++){
nodes[i] = new Node();
nodes[i].count = 1;
nodes[i].sum = input.nextInt();
nodes[i].avg = nodes[i].sum;
res += nodes[i].sum;
}
for(int i = 0 ; i < n - 1; i ++){
int a = input.nextInt();
int b = input.nextInt();
nodes[b].father = a;
}
for(int i = 1 ; i <= n - 1; i++){
//循环n减一次,每次找到等效权值最大的点
double max = -1;
int node = 0;
for(int j = 1 ; j <= n ; j++){
if(j != root && nodes[j].avg > max){
max = nodes[j].avg;
node = j;
}
}
int cur = node;
int father = nodes[cur].father;
//将当前节点和并到父节点
res += nodes[father].count * nodes[cur].sum;
nodes[father].sum += nodes[cur].sum;
nodes[father].count += nodes[cur].count;
nodes[father].avg = 1.0 * nodes[father].sum / nodes[father].count;
nodes[cur].avg = -1;
//将当前节点的子节点合并到父节点下面
for(int j = 1 ; j <= n ; j++){
if(nodes[j].father == cur){
nodes[j].father = father;
}
}
}
System.out.println(res);
}
static class Node{
int father;
int count;
int sum;
double avg;
}
}