Submission #765310

#	Time	Username	Problem	Language	Result	Execution time	Memory
765310		Ronin13	Paths (RMI21_paths)	C++17	100 / 100	231 ms	19108 KiB

Main

#include <bits/stdc++.h>
#define ll long long
#define ull unsigned ll
#define f first
#define s second
#define pii pair<int,int>
#define pll pair<ll,ll>
#define pb push_back
#define epb emplace_back
using namespace std;

const int nmax = 200001;

vector <pair<int, ll> > g[nmax];
ll mx[nmax];

multiset <ll> s, l;
ll ans[nmax];
ll cur;
int n, k;

void add(ll val){
    l.insert(val);
    cur += val;
    if(l.size() > k){
        ll x = *l.begin();
        l.erase(l.begin());
        cur -= x;
        s.insert(x);
    }
}

void rem(ll val){
    if(l.find(val) != l.end()){
        l.erase(l.find(val));
        cur -= val;
        if(l.size() == k) return;
        if(!s.empty()){
            cur += *s.rbegin();
            l.insert(*s.rbegin());
            s.erase(s.find(*s.rbegin()));
        }
    }
    else{
        s.erase(s.find(val));
    }
}

void dfs(int v, int e){

    for(auto& [to, l] : g[v]){
        if(to == e) continue;
        dfs(to, v);
        mx[v] = max(mx[v], mx[to] + l);
    }

    bool x = false;
    for(auto& [to, l] : g[v]){
        if(to == e) continue;
        if(mx[to] + l == mx[v]){
            if(x) add(mx[to] + l);
            else x = true;
        }
        else{
            add(mx[to] + l);
        }
    }
}

void reroot(int v, int e, ll U){
    ll Mx = U, Mx1 = -1e15;

    ans[v] = cur;
    for(auto& [to, l] : g[v]){
        if(to == e) continue;
        if(mx[to] + l > Mx){
            Mx1 = Mx;
            Mx = mx[to] + l;
        }
        else Mx1 = max(Mx1, mx[to] + l);
    }
    for(auto& [to ,l] : g[v]){
        if(to == e) continue;
        rem(mx[to] + l);
        add(mx[to]);
        if(mx[to] + l == Mx){
            rem(Mx1);
            add(Mx1 + l);
            reroot(to, v, Mx1 + l);
        }
        else rem(Mx), add(Mx + l), reroot(to, v, Mx + l);
        if(mx[to] + l == Mx){
            rem(Mx1 + l);
            add(Mx1);
        }
        else rem(Mx + l), add(Mx);
        add(mx[to] + l);
        rem(mx[to]);
    }
}

int main(){
    ios_base::sync_with_stdio(false); cin.tie(0);
    cin >> n;
     cin >> k;
    for(int i = 1; i < n ;i++){
        int u, v, w; cin >> u >> v >> w;
        g[u].pb({v, w});
        g[v].pb({u, w});
    }
    dfs(1, 1);
    //cout << 1;
    add(mx[1]);
    add(0);

    reroot(1, 1, 0);
    for(int i= 1; i <= n ;i++)
        cout << ans[i] << "\n";
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void add(long long int)':
Main.cpp:25:17: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   25 |     if(l.size() > k){
      |        ~~~~~~~~~^~~
Main.cpp: In function 'void rem(long long int)':
Main.cpp:37:21: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   37 |         if(l.size() == k) return;
      |            ~~~~~~~~~^~~~

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