首先可以把题意转化一下:
把一棵点带权树上的 (n) 个点划分为 (A),(B) 两个集合,一种划分方案的代价是:来自 (A) 集合的代价:(sum_{i}d_i+g(A)-sum_{x gets y}[w_x le w_y]);来自 (B) 集合的代价:(sum_{x gets y}[w_x < w_y])。其中 (x gets y) 表示 (x) 是 (y) 的祖先, (g(A)={|A| choose 2}),(d_i) 表示 (i) 到根的距离。总代价为两集合的代价之和。
全部在 (B) 集合的情况非常好算,考虑从 (B) 中逐个取点加入到 (A) 中。
如果点权互不相同,则容易发现取每个点的代价为 (d_i - a_i - b_i),其中 (a_i) 表示比 (i) 小的祖先个数, (b_i) 表示比 (i) 大的子孙个数。它们与 (A,B) 集合的状态无关,所以从小到大排序后一个个加入即可。
如果有点权相同的情况,则代价应该是 (d_i - a_i - b_i - c_i),其中 (c_i) 表示目前 (A) 集合中与 (i) 有祖先子孙关系的点中与 (i) 权值相同的点的个数。这时候我们不知道 (c_i) 是多少,并且不知道取走 (d_i - a_i - b_i - c_i) 最小的点的后续转移是否更优。
不过看起来如果存在祖先子孙权值相同,先取祖先好像更优。因为 (d_x - a_x le d_y - a_y)(尽管可能 (a_x le a_y),但是 (d_x < d_y) 有压倒性优势),并且一定有 (-b_x - c_x le -b_y - c_y),于是当前祖先一定比子孙更优。并且选择祖先的后续转移也比子孙的好,因为选择祖先能“惠及”更多的点有 (-c_i) 的代价。因此,选择一个点的时候与其相等祖先一定全部被选,与其相等的子孙一定全没被选。据此,(c_i) 还可以表示为与 (i) 相等的祖先的个数。
每个点的 (a_i,b_i,c_i,d_i) 可以用树状数组算出。
关键代码:
int fa[N], dep[N];
int a[N], b[N];
void dfs1(int cur, int faa) {
fa[cur] = faa; dep[cur] = dep[faa] + 1;
a[cur] = query(w[cur]);
add(w[cur], 1);
for (register int i = head[cur]; i; i = e[i].nxt) {
int to = e[i].to; if (to == faa) continue;
dfs1(to, cur);
}
add(w[cur], -1);
}
void dfs2(int cur, int faa) {
add(w[cur], 1);
b[cur] -= (query(ltot) - query(w[cur]));
for (register int i = head[cur]; i; i = e[i].nxt) {
int to = e[i].to; if (to == faa) continue;
dfs2(to, cur);
}
b[cur] += (query(ltot) - query(w[cur]));
}
int id[N];
inline bool cmp(const int x, const int y) {
return dep[x] - a[x] - b[x] < dep[y] - a[y] - b[y];
}
ll ans[N];
int main() {
read(n);
for (register int i =1 ; i <= n; ++i) read(w[i]), h[i] = w[i];
lsh();
for (register int i =1 ; i < n; ++i) {
int u, v; read(u), read(v);
addedge(u, v), addedge(v, u);
}
dep[0] = -1;
dfs1(1, 0);
dfs2(1, 0);
for (register int i = 1; i <= n; ++i) id[i] = i;
sort(id +1 , id + 1 + n, cmp);
ll res = 0;
for (register int i = 1; i <= n; ++i) res += b[i];
ans[n] = res;
for (register int i = 1; i <= n; ++i) {
int p = id[i];
res += dep[p] - a[p] - b[p];
ans[n - i] = res + ((1ll * i * (i - 1)) >> 1);
}
for (register int i = 0; i <= n; ++i)
printf("%lld
", ans[i]);
return 0;
}