树上启发式合并

作用

主要解决树上对每个根节点的影响,复杂度最好O(nlog n),最差O(N*N)

#include <bits/stdc++.h>

using namespace std;
#define ll long long
const ll N = 1e5 + 10;
ll dep[N], sze[N], son[N];
ll dian[N];
ll lev[N];
ll max_len;
ll ans[N];
ll flag;
ll he[2 * N], ne[2 * N], to[2 * N];
ll dex;
void add(ll x, ll y) {
    dex++;
    ne[dex] = he[x];
    he[x] = dex;
    to[dex] = y;
}//建边
struct node {
    ll l, r, sum;
} a[N << 4];
//线段树
void pushup(ll x) { a[x].sum = max(a[x << 1].sum, a[x << 1 | 1].sum); }
void build(ll x, ll l, ll r) {
    a[x].l = l, a[x].r = r;
    if (l == r) {
        a[x].sum = lev[l];
        return;
    }
    ll mid = (l + r) >> 1;
    build(x << 1, l, mid);
    build(x << 1 | 1, mid + 1, r);
    pushup(x);
}
void change(ll x, ll u, ll v) {
    if (a[x].l == u && a[x].r == u) {
        a[x].sum += v;
    } else {
        ll mid = a[x].l + a[x].r >> 1;
        if (u <= mid)
            change(x << 1, u, v);
        else
            change(x << 1 | 1, u, v);
        pushup(x);
    }
}
ll query(ll x, ll l, ll r) {
    if (a[x].l >= l && a[x].r <= r)
        return a[x].sum;
    else {
        ll v = 0;
        ll mid = a[x].l + a[x].r >> 1;
        if (l > mid)
            v = query(x << 1 | 1, l, r);
        else if (r <= mid)
            v = query(x << 1, l, r);
        else {
            v = query(x << 1 | 1, l, r);
            v = max(v, query(x << 1, l, r));
        }
        return v;
    }
}
//树链刨分 找重儿子
void dfs_son(int u, int fa) {
    sze[u] = 1, dep[u] = dep[fa] + 1;
    max_len = max(max_len, dep[u]);
    lev[dep[u]] += dian[u];
    ll maxn = 0;
    for (int i = he[u]; i; i = ne[i]) {
        int j = to[i];
        if (j == fa) continue;
        dfs_son(j, u);
        sze[u] += sze[j];
        if (sze[j] > maxn) {
            maxn = sze[j];
            son[u] = j;
        }
    }
}
//计算贡献
void count(int u, int f, int val) {
    change(1, dep[u], dian[u] * (-val));
    for (int i = he[u]; i; i = ne[i]) {
        int v = to[i];
        if (v == f || v == flag) continue;
        count(v, u, val);
    }
}
//树上启发式合并
void dfs(int u, int fa, bool keep) {
    for (int i = he[u]; i; i = ne[i]) {
        int j = to[i];
        if (j == fa || j == son[u]) continue;
        dfs(j, u, 0);
    }
    if (son[u]) {
        dfs(son[u], u, 1);
        flag = son[u];
    }
    count(u, fa, 1);
    ans[u] = query(1, 1, max_len);
    flag = 0;
    if (!keep) {
        count(u, fa, -1);
    }
}
int main() {
    int n;
    cin >> n;
    for (int i = 1; i <= n - 1; i++) {
        int u, v;
        cin >> u >> v;
        add(u, v);
        add(v, u);
    }
    for (int i = 1; i <= n; i++) cin >> dian[i];
    dfs_son(1, 0);
    build(1, 1, max_len);
    dfs(1, 0, 0);
    for (int i = 1; i <= n; i++) cout << ans[i] << " ";
}