Submission #626957

#	Time	Username	Problem	Language	Result	Execution time	Memory
626957		abcvuitunggio	Paths (RMI21_paths)	C++17	100 / 100	261 ms	17952 KiB

Main

#include <iostream>
#include <vector>
#include <set>
#define int long long
using namespace std;
int n,k,u,v,w,dp[100001],dp2[100001],res[100001],x;
vector <pair <int, int> > ke[100001];
multiset <int> s,s2;
void add(int val){
    s.insert(val);
    x+=val;
    if (s.size()>k){
        s2.insert(*s.begin());
        x-=*s.begin();
        s.erase(s.begin());
    }
}
void del(int val){
    if (s2.find(val)!=s2.end()){
        s2.erase(s2.find(val));
        return;
    }
    if (s.find(val)!=s.end()){
        s.erase(s.find(val));
        x-=val;
        s.insert(*--s2.end());
        x+=*--s2.end();
        s2.erase(--s2.end());
    }
}
void dfs(int u, int p){
    for (auto [v,w]:ke[u])
        if (v!=p){
            dfs(v,u);
            dp[u]=max(dp[u],dp[v]+w);
        }
}
void dfs2(int u, int p){
    int mx=-1,mx2=-1,v1=-1;
    for (auto [v,w]:ke[u])
        if (v!=p){
            if (dp[v]+w>mx){
                mx2=mx;
                mx=dp[v]+w;
                v1=v;
            }
            else if (dp[v]+w>mx2)
                mx2=dp[v]+w;
        }
    for (auto [v,w]:ke[u])
        if (v!=p){
            dp2[v]=max(dp2[u],(v==v1?mx2:mx))+w;
            dfs2(v,u);
            if (v!=v1)
                add(dp[v]+w);
        }
}
void dfs3(int u, int p){
    res[u]=x;
    for (auto [v,w]:ke[u])
        if (v!=p){
            del(dp[v]+w);
            add(dp[v]);
            del(dp2[v]-w);
            add(dp2[v]);
            dfs3(v,u);
            add(dp[v]+w);
            del(dp[v]);
            add(dp2[v]-w);
            del(dp2[v]);
        }
}
signed main(){
    ios_base::sync_with_stdio(NULL);cin.tie(nullptr);
    cin >> n >> k;
    for (int i=1;i<n;i++){
        cin >> u >> v >> w;
        ke[u].push_back({v,w});
        ke[v].push_back({u,w});
    }
    dfs(1,1);
    add(dp[1]);
    dfs2(1,1);
    dfs3(1,1);
    for (int i=1;i<=n;i++)
        cout << res[i] << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'void add(long long int)':
Main.cpp:12:17: 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]
   12 |     if (s.size()>k){
      |         ~~~~~~~~^~

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