Submission #1289200

#	Time	Username	Problem	Language	Result	Execution time	Memory
1289200		MisterReaper	Paths (RMI21_paths)	C++20	56 / 100	651 ms	39172 KiB

Main

#include <bits/stdc++.h>

using i64 = long long;

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

std::mt19937 rng(23);

struct node {
    i64 val = 0;
    i64 sum = 0;
    int siz = 1;
    node* lv = nullptr;
    node* rv = nullptr;
    unsigned long int wei = rng();
    node() {}
    node(i64 x) : val(x), sum(x) {}
};

int size(node* v) { return v ? v->siz : 0; }
i64 sum(node* v) { return v ? v->sum : 0; }

void pull(node*& v) {
    if (!v) {
        return;
    }
    v->siz = size(v->lv) + size(v->rv) + 1;
    v->sum = sum(v->lv) + sum(v->rv) + v->val;
}

void merge(node*& v, node* lv, node* rv) {
    if (!lv || !rv) {
        v = (lv ? lv : rv);
        return;
    }
    if (lv->wei > rv->wei) {
        merge(lv->rv, lv->rv, rv);
        v = lv;
    } else {
        merge(rv->lv, lv, rv->lv);
        v = rv;
    }
    pull(v);
}

void split(node* v, node*& lv, node*& rv, int k) {
    if (!v) {
        lv = nullptr;
        rv = nullptr;
        return;
    }
    if (size(v->lv) >= k) {
        split(v->lv, lv, v->lv, k);
        rv = v;
    } else {
        split(v->rv, v->rv, rv, k - size(v->lv) - 1);
        lv = v;
    }
    pull(lv);
    pull(rv);
}

int count(node* v, i64 x) {
    if (!v) {
        return 0;
    }
    if (v->val <= x) {
        return size(v->lv) + 1 + count(v->rv, x);
    } else {
        return count(v->lv, x);
    }
}

node* root = nullptr;

void debug_dfs(node* v) {
    if (!v) {
        return;
    }
    debug_dfs(v->lv);
    std::cerr << v->val << ' ';
    debug_dfs(v->rv);
}

void debug_treap(node* v) {
    #ifdef DEBUG
        std::cerr << "treap: ";
        debug_dfs(v);
        std::cerr << '\n';
    #endif
}

void add(i64 x) {
    debug("add:", x);
    int s = count(root, x);
    node* nd = new node(x);
    node *a, *b;
    split(root, a, b, s);
    merge(root, a, nd);
    merge(root, root, b);
    debug_treap(root);
}

void del(i64 x) {
    debug("del: ", x);
    int s = count(root, x);
    node *a, *b, *c;
    split(root, a, b, s - 1);
    split(b, b, c, 1);
    merge(root, a, c);
    delete b;
    debug_treap(root);
}

int N, K;

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

std::vector<i64> dep;

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;
        del(dep[x]);
        dep[x] += E[i][2];
        add(dep[x]);
        comes[v].emplace(dep[x], x);
    }

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

i64 calc() {
    // auto cd = dep;
    // std::sort(cd.begin(), cd.end(), std::greater());
    // i64 res = 0;
    // for (int i = 0; i < K; ++i) {
    //     res += cd[i];
    // }
    node *a, *b;
    split(root, a, b, N - K);
    i64 res = sum(b);
    merge(root, a, b);
    return res;
}

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;
            }
        }
        del(dep[vtx]);
        del(dep[mxx.second]);
        dep[vtx] -= E[i][2];
        dep[mxx.second] += E[i][2];
        add(dep[vtx]);
        add(dep[mxx.second]);
        ans[u] = calc();
        comes[u].emplace(mxx.first + E[i][2], mxx.second);
        dfs2(u, v);
        del(dep[mxx.second]);
        del(dep[vtx]);
        dep[mxx.second] -= E[i][2];
        dep[vtx] += E[i][2];
        add(dep[mxx.second]);
        add(dep[vtx]);
    }
}

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);

    for (int i = 0; i < N; ++i) {
        add(0);
    }

    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);

    ans[0] = calc();

    dfs2(0, 0);

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

    return 0;
}

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