Submission #1129658

#	Time	Username	Problem	Language	Result	Execution time	Memory
1129658		eysbutno	Paths (RMI21_paths)	C++20	100 / 100	385 ms	18732 KiB

Main

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

using ll = long long;
using pii = array<int, 2>;

#define all(x) begin(x), end(x)
#define sz(x) (int) (x).size()

int main() {
    cin.tie(0) -> sync_with_stdio(0);
    int n, k;
    cin >> n >> k;
    vector<vector<pii>> adj(n);
    for (int i = 1; i < n; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        u--, v--;
        adj[u].push_back({v, w});
        adj[v].push_back({u, w});
    }

    vector<int> leaves;
    vector<array<ll, 2>> best_path(n);
    const auto dfs = [&](int u, int p, auto &&self) -> void {
        best_path[u] = {0, u};
        for (const auto [v, w] : adj[u]) {
            if (v == p) { continue; }
            self(v, u, self);
            best_path[u] = max(best_path[u], {best_path[v][0] + w, best_path[v][1]});
        }
        if (sz(adj[u]) == 0 || (sz(adj[u]) == 1 && adj[u][0][0] == p)) {
            leaves.push_back(u);
        }
    };
    dfs(0, -1, dfs);

    vector<ll> placed(n);
    const auto dfs2 = [&](int u, int p, auto &&self) -> void {
        const auto [_, leaf] = best_path[u];
        if (sz(adj[u]) > 1 || (u == 0)) {
            placed[leaf] += placed[u];
            placed[u] = 0;
        }

        for (const auto [v, w] : adj[u]) {
            if (v == p) { continue; }
            placed[v] += w;
            self(v, u, self);
        }
    };
    dfs2(0, -1, dfs2);
    
    ll sum = 0;
    multiset<ll> big, small;

    const auto fix = [&]() -> void {
        while (sz(big) < k && sz(small)) {
            auto cur = *rbegin(small);
            big.insert(cur);
            small.erase(small.find(cur));
            sum += cur;
        }

        while (sz(small) && sz(big) && *rbegin(small) > *begin(big)) {
            auto val_1 = *rbegin(small);
            auto val_2 = *begin(big);
            small.erase(small.find(val_1));
            big.erase(big.find(val_2));
            small.insert(val_2);
            big.insert(val_1);
            sum += abs(val_1 - val_2);
        }
    };

    const auto ins = [&](ll val) -> void {
        small.insert(val);
        fix();
    };

    const auto del = [&](ll val) -> void {
        if (big.count(val)) { sum -= val; }
        auto &container = (big.count(val)) ? big : small;
        assert(container.count(val) > 0);
        container.erase(container.find(val));
        fix();
    };

    for (int x : leaves) {
        ins(placed[x]);
    }

    vector<int> up(n, -1);
    vector<ll> res(n);
    const auto reroot = [&](int u, int p, auto &&self) -> void {
        res[u] = sum;
        const auto [_, leaf] = best_path[u];
        if (u == 0 && sz(adj[u]) == 1) {
            up[u] = u;
            ins(0);
        }
        ll best_1 = up[u], best_2 = -1;
        for (const auto [v, w] : adj[u]) {
            if (v == p) { continue ;}
            ll leaf_val = placed[best_path[v][1]];
            if (best_1 == -1 || placed[best_1] < leaf_val) {
                best_2 = best_1;
                best_1 = best_path[v][1];
            } else if (best_2 == -1 || placed[best_2] < leaf_val) {
                best_2 = best_path[v][1];
            }
        }

        for (const auto [v, w] : adj[u]) {
            if (v == p) { continue; }
            const int upd = (best_path[v][1] == best_1) ? best_2 : best_1;
            assert(upd != -1);
            up[v] = upd;

            const int nxt = best_path[v][1];
            del(placed[nxt]);
            ins(placed[nxt] -= w);
            del(placed[upd]);
            ins(placed[upd] += w);

            self(v, u, self);

            del(placed[upd]);
            ins(placed[upd] -= w);
            del(placed[nxt]);
            ins(placed[nxt] += w);
        }
    };
    reroot(0, -1, reroot);

    for (int i = 0; i < n; i++) {
        cout << res[i] << "\n";
    }
}

• Subtask 1: 8 / 8
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 20 / 20
• Subtask 5: 12 / 12
• Subtask 6: 32 / 32