HDU 6133 Army Formations

官方题解：BestCoder Blog| 2017 Multi-University Training Contest 8 solutions BY 南京大学

代码：

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

#define fi first
#define se second
#define pb push_back
#define CLR(A, X) memset(A, X, sizeof(A))
#pragma comment(linker, "/STACK:102400000, 102400000")
typedef long long LL;
typedef pair<int, int> PII;
const int N = 1e5+5;

int M, b[N], lson[N], rson[N], sz[N];
LL set_sum, cnt_c[N], sum_c[N], ans[N];
PII a[N];
vector<int> G[N];

int lowbit(int x) { return x&-x; }
void add(LL *c, int x, LL d) { while(x <= M) { c[x] += d; x += lowbit(x); } }
LL sum(LL *c, int x) { LL ret = 0; while(x > 0) { ret += c[x]; x -= lowbit(x); } return ret; }

void Insert(PII x) {
    set_sum += (sum(cnt_c, M)-sum(cnt_c, x.se)+1)*x.fi;
    set_sum += sum(sum_c, x.se);
    add(cnt_c, x.se, 1);
    add(sum_c, x.se, x.fi);
}

void Sub(PII x) {
    add(cnt_c, x.se, -1);
    add(sum_c, x.se, -x.fi);
    set_sum -= (sum(cnt_c, M)-sum(cnt_c, x.se)+1)*x.fi;
    set_sum -= sum(sum_c, x.se);
}

void recover(int u) { if(u == -1) return; Insert(a[u]); recover(lson[u]); recover(rson[u]); }

void Clear(int u) { if(u == -1) return; Sub(a[u]); Clear(lson[u]); Clear(rson[u]); }

void Swap(int u, int fa) {
    lson[u] = rson[u] = -1;
    sz[u] = 1;
    for(int v:G[u]) if(v != fa) {
        if(~rson[u]) lson[u] = v;
        else rson[u] = v;
        Swap(v, u);
        sz[u] += sz[v];
    }
    if(~lson[u] && ~rson[u] && sz[lson[u]]>sz[rson[u]]) swap(lson[u], rson[u]);
}

void solve(int u) {
    if(~lson[u]) solve(lson[u]);
    if(~rson[u]) {
        Clear(lson[u]);
        solve(rson[u]);
        recover(lson[u]);
    }
    Insert(a[u]);
    ans[u] = set_sum;
}

void init(int n) {
    set_sum = 0;
    for(int i = 1; i <= n; i++) G[i].clear();
    sort(b+1, b+n+1);
    M = unique(b+1, b+n+1)-(b+1);
    for(int i = 1; i <= n; i++) {
        a[i].se = lower_bound(b+1, b+M+1, a[i].fi)-b;
        cnt_c[i] = sum_c[i] = 0;
    }
}

int main() {
    int T;
    scanf("%d", &T);
    while(T--) {
        int n, u, v;
        scanf("%d", &n);
        for(int i = 1; i <= n; i++) {
            scanf("%d", &b[i]);
            a[i].fi = b[i];
        }
        init(n);
        for(int i = 1; i < n; i++) {
            scanf("%d%d", &u, &v);
            G[u].pb(v); G[v].pb(u);
        }
        Swap(1, -1);
        solve(1);
        for(int i = 1; i <= n; i++) printf("%lld ", ans[i]);
        puts("");
    }
    return 0;
}