Submission #1289219

#TimeUsernameProblemLanguageResultExecution timeMemory
1289219MisterReaperPaths (RMI21_paths)C++20
100 / 100
162 ms42100 KiB
#include <bits/stdc++.h>

using i64 = long long;

#ifdef DEBUG
    #include "debug.h"
#else
    #define debug(...) void(23)
#endif

int N, K;

std::vector<std::vector<int>> adj;
std::vector<std::array<int, 3>> E;

std::vector<i64> dep;

std::vector<int> col;
std::vector<std::set<std::pair<i64, int>, std::greater<std::pair<i64, int>>>> comes;
std::vector<i64> ans;
std::vector<int> saver;

int dfs1(int v, int pr) {
    debug(v, pr);
    if (int(adj[v].size()) == 1) {
        comes[v].emplace(0, v);
    }
    for (auto i : adj[v]) {
        int u = E[i][0] ^ E[i][1] ^ v;
        if (u == pr) {
            continue;
        }
        int x = dfs1(u, v);
        saver[u] = x;
        dep[x] += E[i][2];
        comes[v].emplace(dep[x], x);
    }

    return comes[v].begin()->second;
}

int siz = 0;
i64 sum = 0;
std::priority_queue<std::pair<i64, int>, std::vector<std::pair<i64, int>>, std::greater<std::pair<i64, int>>> pq1;
std::priority_queue<std::pair<i64, int>, std::vector<std::pair<i64, int>>> pq2;

void fix() {
    debug(siz);
    while (siz < K) {
        while (!pq2.empty() && (dep[pq2.top().second] != pq2.top().first || col[pq2.top().second] == true)) {
            pq2.pop();
        }
        siz += 1;
        sum += dep[pq2.top().second];
        col[pq2.top().second] = true;
        pq1.emplace(pq2.top());
        pq2.pop();
    }
    debug(siz);
    while (true) {
        while (!pq1.empty() && (dep[pq1.top().second] != pq1.top().first || col[pq1.top().second] == false)) {
            pq1.pop();
        }
        while (!pq2.empty() && (dep[pq2.top().second] != pq2.top().first || col[pq2.top().second] == true)) {
            pq2.pop();
        }
        if (pq2.empty() || pq1.top().first >= pq2.top().first) {
            break;
        }
        sum += pq2.top().first - pq1.top().first;
        auto[a, b] = pq1.top();
        auto[c, d] = pq2.top();
        pq1.pop();
        pq2.pop();
        col[b] = false;
        col[d] = true;
        pq1.emplace(c, d);
        pq2.emplace(a, b);
    }
}

void dfs2(int v, int pr) {
    debug(v, pr, dep);
    for (auto i : adj[v]) {
        int u = E[i][0] ^ E[i][1] ^ v;
        if (u == pr) {
            continue;
        }   
        int vtx = saver[u];
        std::pair<i64, int> mxx = {-1, -1};
        for (auto[d, vv] : comes[v]) {
            if (vv != vtx) {
                mxx = {d, vv};
                break;
            }
        }
        
        if (col[vtx]) {
            sum -= dep[vtx];
            siz -= 1;
            pq2.emplace(dep[vtx] - E[i][2], vtx);
            col[vtx] = false;
        } else {
            pq2.emplace(dep[vtx] - E[i][2], vtx);
        }
        dep[vtx] -= E[i][2];

        auto xu = mxx.second;
        if (col[xu]) {
            sum -= dep[xu];
            siz -= 1;
            pq2.emplace(dep[xu] + E[i][2], xu);
            col[xu] = false;
        } else {
            pq2.emplace(dep[xu] + E[i][2], xu);
        }
        dep[xu] += E[i][2];
        
        debug(1);
        fix();
        debug(2);
        ans[u] = sum;
        comes[u].emplace(dep[xu], xu);
        dfs2(u, v);

        if (col[xu]) {
            sum -= dep[xu];
            siz -= 1;
            pq2.emplace(dep[xu] - E[i][2], xu);
            col[xu] = false;
        } else {
            pq2.emplace(dep[xu] - E[i][2], xu);
        }
        dep[xu] -= E[i][2];

        if (col[vtx]) {
            sum -= dep[vtx];
            siz -= 1;
            pq2.emplace(dep[vtx] + E[i][2], vtx);
            col[vtx] = false;
        } else {
            pq2.emplace(dep[vtx] + E[i][2], vtx);
        }
        dep[vtx] += E[i][2];
    }
}

int main() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(nullptr);

    std::cin >> N >> K;

    adj.resize(N);
    E.resize(N);
    comes.resize(N);
    ans.resize(N);
    saver.resize(N);
    dep.resize(N);
    col.resize(N);

    for (int i = 1; i < N; ++i) {
        int A, B, C;
        std::cin >> A >> B >> C;
        --A, --B;
        E[i] = {A, B, C};
        adj[A].emplace_back(i);
        adj[B].emplace_back(i);
    }

    dfs1(0, 0);

    std::vector<int> ord(N);
    std::iota(ord.begin(), ord.end(), 0);
    std::sort(ord.begin(), ord.end(), [&](auto lhs, auto rhs) {
        return dep[lhs] > dep[rhs];
    });
    for (int i = 0; i < K; ++i) {
        siz += 1;
        ans[0] += dep[ord[i]];
        sum += dep[ord[i]];
        pq1.emplace(dep[ord[i]], ord[i]);
        col[ord[i]] = true;
    }
    for (int i = K; i < N; ++i) {
        pq2.emplace(dep[ord[i]], ord[i]);
    }

    sum = ans[0];

    dfs2(0, 0);

    for (int i = 0; i < N; ++i) {
        std::cout << ans[i] << '\n';
    }

    return 0;
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...