Submission #525113

#	Time	Username	Problem	Language	Result	Execution time	Memory
525113		two_sides	Paths (RMI21_paths)	C++17	100 / 100	195 ms	22508 KiB

Main

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

const int N = 100005;

ll dp[N], ans[N]; int root, n, k;
vector<pair<int, int>> adj[N];
multiset<ll> big, small;
ll sum_big, sum_small;

void erase(ll x) {
    if (small.count(x)) {
        small.erase(small.find(x));
        sum_small -= x;
    } else {
        big.erase(big.find(x));
        sum_big -= x;
    }
}

void insert(ll x) {
    if (big.empty() || x < *big.begin()) {
        small.insert(x);
        sum_small += x;
    } else {
        big.insert(x);
        sum_big += x;
    }
}

void dfs1(int u) {
    int best = 0;
    for (auto e : adj[u]) {
        int v, w; tie(v, w) = e;
        adj[v].erase(find(adj[v].begin(),
        adj[v].end(), make_pair(u, w)));
        dfs1(v);
        if (dp[u] < dp[v] + w) {
            dp[u] = dp[v] + w; best = v;
        }
    }
    for (auto e : adj[u]) {
        int v, w; tie(v, w) = e;
        if (v != best) insert(dp[v] + w);
    }
}

void adjust() {
    while (big.size() > k) {
        auto it = big.begin();
        sum_big -= *it;
        sum_small += *it;
        big.erase(it);
        small.insert(*it);
    }
    while (small.size() && big.size() < k) {
        auto it = prev(small.end());
        sum_small -= *it;
        sum_big += *it;
        small.erase(it);
        big.insert(*it);
    }
}

void dfs2(int u, ll up) {
    adjust(); ans[u] = sum_big;
    vector<ll> pre(1), suf(1);
    for (int i = 0; i < adj[u].size(); i++) {
        int v, w; tie(v, w) = adj[u][i];
        pre.push_back(max(pre.back(), dp[v] + w));
    }
    for (int i = adj[u].size() - 1; i >= 0; i--) {
        int v, w; tie(v, w) = adj[u][i];
        suf.push_back(max(suf.back(), dp[v] + w));
    }
    reverse(suf.begin(), suf.end());
    for (int i = 0; i < adj[u].size(); i++) {
        int v, w; tie(v, w) = adj[u][i];
        ll best = max({up, pre[i], suf[i + 1]});
        erase(best); insert(best + w);
        erase(dp[v] + w); insert(dp[v]);
        dfs2(v, best + w);
        erase(best + w); insert(best);
        erase(dp[v]); insert(dp[v] + w);
    }
}

int main() {
    cin.tie(0)->sync_with_stdio(0);
    cin >> n >> k;
    if (n == 1) {
        cout << "0\n"; return 0;
    }
    if (n == 2) {
        int u, v, w;
        cin >> u >> v >> w;
        cout << w << '\n' << w << '\n';
        return 0;
    }
    for (int i = 1; i < n; i++) {
        int u, v, w;
        cin >> u >> v >> w;
        adj[u].emplace_back(v, w);
        adj[v].emplace_back(u, w);
    }
    for (int i = 1; i <= n; i++)
        if (adj[i].size() > 1)
            root = i;
    dfs1(root); insert(dp[root]);
    dfs2(root, 0);
    for (int i = 1; i <= n; i++)
        cout << ans[i] << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'void adjust()':
Main.cpp:52:23: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   52 |     while (big.size() > k) {
      |            ~~~~~~~~~~~^~~
Main.cpp:59:39: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   59 |     while (small.size() && big.size() < k) {
      |                            ~~~~~~~~~~~^~~
Main.cpp: In function 'void dfs2(int, ll)':
Main.cpp:71:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   71 |     for (int i = 0; i < adj[u].size(); i++) {
      |                     ~~^~~~~~~~~~~~~~~
Main.cpp:80:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   80 |     for (int i = 0; i < adj[u].size(); i++) {
      |                     ~~^~~~~~~~~~~~~~~

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