Submission #1129650

#	Time	Username	Problem	Language	Result	Execution time	Memory
1129650		eysbutno	Paths (RMI21_paths)	C++20	0 / 100	59 ms	10564 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--;
        if (u > v) { swap(u, v); }
        adj[u].push_back({v, w});
    }

    vector<array<ll, 2>> best_path(n);
    
    // calculates down and best_path
    for (int u = n - 1; u >= 0; u--) {
        best_path[u] = {0, u};
        for (const auto [v, w] : adj[u]) {
            best_path[u] = max(best_path[u], {best_path[v][0] + w, best_path[v][1]});
        }
    }

    vector<ll> placed(n);
    for (int u = 0; u < n; u++) {
        const auto [_, leaf] = best_path[u];
        if (sz(adj[u])) {
            placed[leaf] += placed[u];
            placed[u] = 0;
        }

        for (const auto [v, w] : adj[u]) {
            placed[v] += w;
        }
    }

    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 i = 0; i < n; i++) {
        if (sz(adj[i]) == 0) { ins(placed[i]); }
    }

    vector<int> up(n, -1);
    vector<ll> res(n);
    const auto reroot = [&](int u, 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]) {
            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];
            }
        }

        // if (u == 0) { cout << "asdf " << best_1 << " " << best_2 << "\n"; }
        // cout << u << " " << up[u] << "\n";

        for (const auto [v, w] : adj[u]) {
            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, self);

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

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

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