Dynamic Values And Maximum Sum

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

using ll = long long;

void setIO() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
}

struct Best {
    int leaf;
    int dist;
};

const int NEG = -1000000000;

bool betterBest(const Best &a, const Best &b) {
    if (a.dist != b.dist) return a.dist > b.dist;
    if (a.leaf == 0) return false;
    if (b.leaf == 0) return true;
    return a.leaf < b.leaf;
}

void addBest(Best x, Best &first, Best &second) {
    if (betterBest(x, first)) {
        second = first;
        first = x;
    } else if (betterBest(x, second)) {
        second = x;
    }
}

struct TopK {
    int need = 0;
    vector<ll> *val = nullptr;
    vector<char> active;
    set<pair<ll, int>> top, rest;
    ll sum = 0;

    TopK() = default;

    TopK(int n, int k, vector<ll> &bucket) {
        need = k;
        val = &bucket;
        active.assign(n + 1, 0);
    }

    void rebalance() {
        while ((int)top.size() > need) {
            auto it = top.begin();
            auto p = *it;
            top.erase(it);
            sum -= p.first;
            rest.insert(p);
        }
        while ((int)top.size() < need && !rest.empty()) {
            auto it = prev(rest.end());
            auto p = *it;
            rest.erase(it);
            top.insert(p);
            sum += p.first;
        }
        while (!top.empty() && !rest.empty()) {
            auto itSmall = top.begin();
            auto itBig = prev(rest.end());
            if (!(*itSmall < *itBig)) break;

            auto small = *itSmall;
            auto big = *itBig;
            top.erase(itSmall);
            rest.erase(itBig);

            sum += big.first - small.first;
            top.insert(big);
            rest.insert(small);
        }
    }

    void insertActivePair(int x) {
        pair<ll, int> p = {(*val)[x], x};
        top.insert(p);
        sum += p.first;
        rebalance();
    }

    void eraseActivePair(int x) {
        pair<ll, int> p = {(*val)[x], x};
        auto it = top.find(p);
        if (it != top.end()) {
            sum -= p.first;
            top.erase(it);
        } else {
            it = rest.find(p);
            if (it != rest.end()) rest.erase(it);
        }
        rebalance();
    }

    void addCandidate(int x) {
        active[x] = 1;
        insertActivePair(x);
    }

    void removeCandidate(int x) {
        eraseActivePair(x);
        active[x] = 0;
    }

    void changeBucket(int x, ll delta) {
        if (x == 0) return;
        if (active[x]) eraseActivePair(x);
        (*val)[x] += delta;
        if (active[x]) insertActivePair(x);
    }
};

int main() {
    setIO();

    int T;
    cin >> T;
    while (T--) {
        int n, k;
        cin >> n >> k;

        vector<ll> a(n + 1);
        ll maxA = 0;
        for (int i = 1; i <= n; i++) {
            cin >> a[i];
            maxA = max(maxA, a[i]);
        }

        vector<vector<int>> g(n + 1);
        for (int i = 0; i < n - 1; i++) {
            int u, v;
            cin >> u >> v;
            g[u].push_back(v);
            g[v].push_back(u);
        }

        if (k == 1) {
            cout << maxA << '\n';
            continue;
        }

        vector<int> parent(n + 1, 0), order;
        order.reserve(n);
        parent[1] = -1;
        vector<int> st = {1};
        while (!st.empty()) {
            int v = st.back();
            st.pop_back();
            order.push_back(v);
            for (int to : g[v]) {
                if (to == parent[v]) continue;
                parent[to] = v;
                st.push_back(to);
            }
        }

        vector<Best> down1(n + 1), down2(n + 1), up(n + 1);
        for (int idx = n - 1; idx >= 0; idx--) {
            int v = order[idx];
            down1[v] = {v, 0};
            down2[v] = {0, NEG};
            for (int to : g[v]) {
                if (parent[to] != v) continue;
                Best cand = {down1[to].leaf, down1[to].dist + 1};
                addBest(cand, down1[v], down2[v]);
            }
        }

        vector<int> parentDest(n + 1, 0);
        up[1] = {0, NEG};
        for (int v : order) {
            for (int to : g[v]) {
                if (parent[to] != v) continue;

                Best inside = (down1[v].leaf == down1[to].leaf ? down2[v] : down1[v]);
                Best best = inside;
                if (betterBest(up[v], best)) best = up[v];

                parentDest[to] = best.leaf;
                up[to] = {best.leaf, best.dist + 1};
            }
        }

        vector<ll> bucket(n + 1, 0);
        vector<int> dest(n + 1, 0);
        for (int v = 2; v <= n; v++) {
            dest[v] = down1[v].leaf;
            bucket[dest[v]] += a[v];
        }

        TopK ds(n, k - 1, bucket);
        for (int v = 2; v <= n; v++) ds.addCandidate(v);

        ll answer = 0;

        auto moveRoot = [&](int p, int v) {
            ds.changeBucket(dest[v], -a[v]);
            dest[v] = 0;

            ds.removeCandidate(v);
            ds.addCandidate(p);

            dest[p] = parentDest[v];
            ds.changeBucket(dest[p], a[p]);
        };

        auto undoMove = [&](int p, int v) {
            ds.changeBucket(dest[p], -a[p]);
            dest[p] = 0;

            ds.removeCandidate(p);
            ds.addCandidate(v);

            dest[v] = down1[v].leaf;
            ds.changeBucket(dest[v], a[v]);
        };

        struct Event {
            int type;
            int p;
            int v;
        };

        vector<Event> events;
        events.push_back({0, 0, 1});
        while (!events.empty()) {
            Event cur = events.back();
            events.pop_back();

            if (cur.type == 0) {
                int v = cur.v;
                answer = max(answer, a[v] + ds.sum);
                for (int to : g[v]) {
                    if (parent[to] == v) events.push_back({1, v, to});
                }
            } else if (cur.type == 1) {
                moveRoot(cur.p, cur.v);
                events.push_back({2, cur.p, cur.v});
                events.push_back({0, cur.p, cur.v});
            } else {
                undoMove(cur.p, cur.v);
            }
        }

        cout << answer << '\n';
    }

    return 0;
}