Submission #538714

#	Time	Username	Problem	Language	Result	Execution time	Memory
538714		siewjh	Paths (RMI21_paths)	C++17	100 / 100	345 ms	24172 KiB

Main

#include <bits/stdc++.h>
using namespace std;
const int MAXN = 100'005;
typedef long long ll;
vector<pair<int, ll>> adjlist[MAXN];
ll maxof[MAXN], ans[MAXN];
int nodes, choose;
ll curr = 0;
multiset<ll> cset, rejset;
void add(ll x){
    curr += x;
    cset.insert(x);
    if (cset.size() > choose){
        auto av = cset.begin();
        curr -= *av;
        rejset.insert(*av);
        cset.erase(av);
    }
}
void rem(ll x){
    if (cset.find(x) == cset.end()) rejset.erase(rejset.find(x));
    else{
        curr -= x;
        cset.erase(cset.find(x));
        if (!rejset.empty()){
            auto av = prev(rejset.end());
            curr += *av;
            cset.insert(*av);
            rejset.erase(av);
        }
    }
}
void dfs(int x, int par){
    for (auto nxt : adjlist[x]){
        int nn = nxt.first; ll nd = nxt.second;
        if (nn == par) continue;
        dfs(nn, x);
        rem(maxof[nn]); add(maxof[nn] + nd);
        maxof[x] = max(maxof[x], maxof[nn] + nd);
    }
    add(0);
}
void dfs2(int x, int par){
    ans[x] = curr;
    vector<pair<ll, int>> comp;
    comp.push_back({0, x});
    for (auto nxt : adjlist[x]) comp.push_back({maxof[nxt.first] + nxt.second, nxt.first});
    sort(comp.begin(), comp.end(), greater<pair<ll, int>>());
    auto f = comp[0], s = comp[1];
    for (auto nxt : adjlist[x]){
        int nn = nxt.first; ll nd = nxt.second;
        if (nn == par) continue;
        ll holder = maxof[x];
        if (f.second == nn) maxof[x] = s.first;
		else maxof[x] = f.first;
		rem(maxof[nn] + nd); add(maxof[nn]);
		rem(maxof[x]); add(maxof[x] + nd);
		dfs2(nn, x);
		rem(maxof[x] + nd); add(maxof[x]);
		rem(maxof[nn]); add(maxof[nn] + nd);
		maxof[x] = holder;
    }
}
int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);
    cin >> nodes >> choose;
    for (int i = 1; i < nodes; i++){
        int a, b; ll c; cin >> a >> b >> c;
        adjlist[a].push_back({b, c});
        adjlist[b].push_back({a, c});
    }
    dfs(1, -1);
    dfs2(1, -1);
    for (int i = 1; i <= nodes; i++) cout << ans[i] << '\n';
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void add(ll)':
Main.cpp:13:21: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   13 |     if (cset.size() > choose){
      |         ~~~~~~~~~~~~^~~~~~~~

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