山东理工大学第十二届ACM程序设计竞赛 - Cut the tree 树上启发式合并

Cut the tree

题意

一棵 $n(2\le n\le 1\cdot 10^5)$ 个节点的树，$1$ 号结点是根节点，每个点有一个点权。树的第 $H$ 层权值为深度为 $H$ 的所有点权之和。树的总权值为所有层权值的最大值。问分别割掉以 $1,2,\dots ,n$ 为根的子树后，剩余树的权值是多少。

解法

使用线段树维护每一层的权值的区间最大值，子树问题考虑树上启发式合并，每次暴力修改，然后还原轻孩子的影响即可。复杂度为 $O(nlo{g}^{2}n)$ 。

代码

#pragma region
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <map>
#include <queue>
#include <set>
#include <vector>
using namespace std;
typedef long long ll;
#define tr t[root]
#define lson t[root << 1]
#define rson t[root << 1 | 1]
#define rep(i, a, n) for (int i = a; i <= n; ++i)
#define per(i, a, n) for (int i = n; i >= a; --i)
#pragma endregion
const int maxn = 1e5 + 5;
int n;
vector<int> g[maxn], d[maxn];
int sz[maxn], son[maxn], dep[maxn];
ll a[maxn], dsum[maxn];
int flag;
ll ans[maxn];
void dfs1(int u, int f) {
    sz[u] = 1, dep[u] = dep[f] + 1;
    d[dep[u]].push_back(u);
    dsum[dep[u]] += a[u];
    for (auto v : g[u]) {
        if (v == f) continue;
        dfs1(v, u);
        sz[u] += sz[v];
        if (sz[v] > sz[son[u]]) son[u] = v;
    }
}
struct segtree {
    int l, r;
    ll maxx;
} t[maxn << 2];
void build(int root, int l, int r) {
    tr.l = l, tr.r = r;
    if (l == r) {
        tr.maxx = dsum[l];
        return;
    }
    int mid = (l + r) >> 1;
    build(root << 1, l, mid);
    build(root << 1 | 1, mid + 1, r);
    tr.maxx = max(lson.maxx, rson.maxx);
}
ll query(int root, int l, int r) {
    if (l <= tr.l && tr.r <= r) return tr.maxx;
    ll ans = 0;
    int mid = (tr.l + tr.r) >> 1;
    if (l <= mid) ans = max(ans, query(root << 1, l, r));
    if (r > mid) ans = max(ans, query(root << 1 | 1, l, r));
    return ans;
}
void update(int root, int pos, ll x) {
    if (tr.l == tr.r) {
        tr.maxx -= x;
        return;
    }
    int mid = (tr.l + tr.r) >> 1;
    if (pos <= mid) update(root << 1, pos, x);
    if (pos > mid) update(root << 1 | 1, pos, x);
    tr.maxx = max(lson.maxx, rson.maxx);
}
void count(int u, int f, int val) {
    update(1, dep[u], a[u] * val);
    for (auto v : g[u]) {
        if (v == f || v == flag) continue;
        count(v, u, val);
    }
}
void dfs(int u, int f, bool keep) {
    for (auto v : g[u]) {
        if (v == f || v == son[u]) continue;
        dfs(v, u, 0);
    }
    if (son[u]) {
        dfs(son[u], u, 1);
        flag = son[u];
    }
    count(u, f, 1);
    ans[u] = query(1, 1, n);
    flag = 0;
    if (!keep) {
        count(u, f, -1);
    }
}
int main() {
    scanf("%d", &n);
    rep(i, 1, n - 1) {
        int u, v;
        scanf("%d%d", &u, &v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    rep(i, 1, n) scanf("%lld", &a[i]);
    dfs1(1, 0);
    build(1, 1, n);
    dfs(1, 0, 0);
    rep(i, 1, n) printf("%lld%c", ans[i], " \n"[i == n]);
}