贪心 + 增量法

需要思考的地方非常多, 很难小篇幅说情况, 详情查看大雪菜的题解或者视频.

时间复杂度

O(N^2)
本来想试图写成O(N*log(N))的, 即, 通过大根堆来寻找均值最大的节点, 但是发现更新fa节点的均值后, 很难在大根堆里面也相应进行更新, 就放弃了.

Python 代码

from collections import defaultdict

class Node:
    def __init__(self):
        self.sum = 0
        self.cnt = 1
        self.avg = 0

n, root = map(int, input().split())
A = list(map(int, input().split()))
nodes = defaultdict(Node)
ans = 0
for i in range(1, n + 1):
    nodes[i].sum = nodes[i].avg = A[i - 1]
    ans += A[i - 1]

u2v = defaultdict(list)
v2u = defaultdict(int)
for i in range(n - 1):
    a, b = map(int, input().split())
    u2v[a].append(b)
    v2u[b] = a

for i in range(n - 1):
    avg_max = -1
    idx_max = -1
    for idx in range(1, n + 1):
        if idx != root and nodes[idx].avg > avg_max:
            avg_max = nodes[idx].avg
            idx_max = idx
    fa = v2u[idx_max]
    ans += nodes[fa].cnt * nodes[idx_max].sum
    nodes[fa].sum += nodes[idx_max].sum
    nodes[fa].cnt += nodes[idx_max].cnt
    nodes[fa].avg = nodes[fa].sum / nodes[fa].cnt

    nodes[idx_max].avg = -1
    u2v[fa].remove(idx_max)
    for child in u2v[idx_max]:
        v2u[child] = fa
        u2v[fa].append(child)
print(ans)