Submission #871280

#	Time	Username	Problem	Language	Result	Execution time	Memory
871280		Ahmed57	Paths (RMI21_paths)	C++17	100 / 100	250 ms	24172 KiB

Main

#include <bits/stdc++.h>
using namespace std;
#define int long long
pair<long long,long long> vl[100001];
long long ans[100001];set<int> leafs;
vector<pair<int,long long>> adj[100001];
void dfs(int i,int pr,int la){
    if(adj[i].size()==1)leafs.insert(i);
    if(adj[i].size()==1&&i!=1){
        ans[i] = la;vl[i] = {la,i};
        return ;
    }
    for(auto j:adj[i]){
        if(j.first==pr)continue;
        dfs(j.first,i,j.second);
        vl[i] = max(vl[i],vl[j.first]);
    }
    ans[vl[i].second]+=la;
    vl[i].first+=la;
}multiset<long long> used,unused;
long long all = 0;int n,k;
void add(long long x){
    all+=x;
    used.insert(x);
    if(used.size()>k){
        all-=*used.begin();
        unused.insert(*used.begin());
        used.erase(used.begin());
    }
}
void rem(long long x){
    if(x<(*used.begin())){
        unused.erase(unused.lower_bound(x));
    }else{
        all-=x;
        used.erase(used.lower_bound(x));
        while(used.size()<k&&unused.size()>0){
            auto it = unused.end();it--;
            used.insert(*it);
            all+=(*it);
            unused.erase(it);
        }
    }
}
long long v[100001];
pair<long long,long long> pref[100001],suf[100001];
void reroot(int i,int pr,pair<long long,long long> pa){
    v[i] = all;
    if(adj[i].size()==1&&i!=1)return;
    if(n==1&&i==1)return;
    if(i==1&&adj[i].size()==1){
        pa = {0,i};
    }
    pair<long long,long long> ma = pa;
    for(auto j:adj[i]){
        if(j.first==pr)continue;
        pref[j.first] = ma;
        ma = max(ma,make_pair(vl[j.first].first,vl[j.first].second));
    }
    ma = pa;
    reverse(adj[i].begin(),adj[i].end());
    for(auto j:adj[i]){
        if(j.first==pr)continue;
        suf[j.first] = ma;
        ma = max(ma,make_pair(vl[j.first].first,vl[j.first].second));
    }
    reverse(adj[i].begin(),adj[i].end());
    for(auto j:adj[i]){
        if(j.first==pr)continue;
        rem(ans[vl[j.first].second]);
        ans[vl[j.first].second]-=j.second;
        add(ans[vl[j.first].second]);
        pair<int,int> xd = max(pref[j.first],suf[j.first]);
        rem(ans[xd.second]);
        ans[xd.second]+=j.second;
        add(ans[xd.second]);
        xd.first+=j.second;
        reroot(j.first,i,xd);
        rem(ans[xd.second]);
        ans[xd.second]-=j.second;
        add(ans[xd.second]);
        rem(ans[vl[j.first].second]);
        ans[vl[j.first].second]+=j.second;
        add(ans[vl[j.first].second]);
    }
}
signed main(){
    cin>>n>>k;
    for(int i =0;i<n-1;i++){
        int a,b,c;cin>>a>>b>>c;
        adj[a].push_back({b,c});
        adj[b].push_back({a,c});
    }
    dfs(1,0,0);
    for(auto i:leafs){
        add(ans[i]);
    }
    reroot(1,0,{0,0});
    for(int i = 1;i<=n;i++){
        cout<<v[i]<<"\n";
    }
}
/*
11 3
1 2 5
2 3 3
2 6 5
3 4 4
3 5 2
1 7 6
7 8 4
7 9 5
1 10 1
10 11 1
*/

Compilation message (stderr)

Main.cpp: In function 'void add(long long int)':
Main.cpp:25:19: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   25 |     if(used.size()>k){
      |        ~~~~~~~~~~~^~
Main.cpp: In function 'void rem(long long int)':
Main.cpp:37:26: warning: comparison of integer expressions of different signedness: 'std::multiset<long long int>::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
   37 |         while(used.size()<k&&unused.size()>0){
      |               ~~~~~~~~~~~^~

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