Submission #647474

#	Time	Username	Problem	Language	Result	Execution time	Memory
647474		k_balint31415	Paths (RMI21_paths)	C++14	100 / 100	318 ms	20992 KiB

Main

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
const int c=1e5+5;


int n,k;

multiset<ll> nagy,kicsi;
ll ossz=0;

void ins(ll a){
    if(nagy.size()<k) {ossz+=a; nagy.insert(a); return;}
    ll nb=*nagy.begin();
    if(nb>=a) {kicsi.insert(a); return;}
    kicsi.insert(nb);
    ossz-=nb;
    nagy.erase(nagy.find(nb));
    ossz+=a; nagy.insert(a);
}

void rem(ll a){
    if(kicsi.find(a) != kicsi.end()){
        kicsi.erase(kicsi.find(a)); return;
    }
    if(nagy.find(a) == nagy.end()) return;
    ossz-=a;
    nagy.erase(nagy.find(a));
    if(!kicsi.empty()) {auto it=kicsi.end(); it--; ll b=*it; ossz+=b; nagy.insert(b); kicsi.erase(kicsi.find(b));}
}

vector<pair<int,ll>> adj[c];
ll dp[c][2];
ll ans[c];

ll dfs(int v, int p, ll len){
    ll maxi=0;
    for(pair<int,ll> x: adj[v]){
        if(x.first != p){
            ll cur=dfs(x.first,v,x.second);
            if(cur>dp[v][1]) dp[v][1]=cur;
            if(dp[v][1]>dp[v][0]) swap(dp[v][1],dp[v][0]);
            if(cur>maxi){
                swap(cur,maxi);
            }
            ins(cur);
        }
    }
    return maxi+len;
}

void dfs2(int v, int p){
    ans[v]=ossz;
    for(pair<int,ll> x: adj[v]){
        if(x.first != p){
            ll A=dp[x.first][0],B=dp[x.first][1];

            rem(dp[x.first][0]+x.second);
            ins(dp[x.first][0]);
            
            bool b=dp[v][0]==(dp[x.first][0]+x.second);
            rem(!b?dp[v][0]:dp[v][1]);
            ins((!b?dp[v][0]:dp[v][1])+x.second);

            dp[x.first][1]=max((!b?dp[v][0]:dp[v][1])+x.second,dp[x.first][1]);
            if(dp[x.first][0]<dp[x.first][1]) swap(dp[x.first][0],dp[x.first][1]);           

            dfs2(x.first,v);
            dp[x.first][0]=A; dp[x.first][1]=B;

            rem((!b?dp[v][0]:dp[v][1])+x.second);
            ins(!b?dp[v][0]:dp[v][1]);

            rem(dp[x.first][0]);
            ins(dp[x.first][0]+x.second);
        }
    }
}

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

    cin>>n>>k;
    for(int i=1;i<n;i++){
        int a,b,val;
        cin>>a>>b>>val;
        adj[a].push_back(make_pair(b,val));
        adj[b].push_back(make_pair(a,val));
    }

    ins(dfs(1,0,0));
    dfs2(1,0);

    for(int i=1;i<=n;i++){
        cout << ans[i] << '\n';
    }
}

Compilation message (stderr)

Main.cpp: In function 'void ins(ll)':
Main.cpp:13:19: 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(nagy.size()<k) {ossz+=a; nagy.insert(a); 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