Submission #1020500

#	Time	Username	Problem	Language	Result	Execution time	Memory
1020500		NintsiChkhaidze	Paths (RMI21_paths)	C++17	100 / 100	335 ms	17620 KiB

Main

#include <bits/stdc++.h>
#include <ostream>
#define pb push_back
#define s second
#define f first
#define pb push_back
#define pii pair <int,int>
#define ll long long
#define left h*2,l,(l + r)/2
#define right h*2+1,(l + r)/2 + 1,r
#define int ll
using namespace std;
 
const int N = 1e5 + 5,inf = 1e18;
vector <pii> v[N];
int n,k,val[N],total,ans[N];
multiset <int> s1,s2;

//s1,s2
void upd(int x,int dl){
  if (dl == 1){
    s1.insert(x);
    total += x;
  }else {
    if (s2.count(x)) s2.erase(s2.find(x));
    else if (s1.count(x)) s1.erase(s1.find(x)),total -= x;
    else assert(0); ///fail
  }

  while (s1.size() < k && s2.size()){
    int x = *(--s2.end());
    s1.insert(x);
    total += x;
    s2.erase(--s2.end());
  }

  while (s1.size() > k){
    int x = *s1.begin();
    s2.insert(x);
    s1.erase(s1.begin());
    total -= x;
  }
}

void dfs(int x,int par){

  for (auto [to,w]: v[x]){
    if (to == par) continue;
    dfs(to,x);
    val[x] = max(val[x],val[to] + w);
  }

  bool foundmax=0;
  for (auto [to,w]: v[x]){
    if (to == par) continue;
    if ((!foundmax) && val[to] + w == val[x]){
      foundmax = 1;
    }else{
      upd(val[to] + w,1);
    }
  }
}

void DFS(int x,int par,int mx){
  ans[x] = total;
  int m1=mx,m2=-1;
  for (auto [to,w]: v[x]){
    if (to==par) continue;
    if (val[to] + w >= m1) m2 = m1,m1 = val[to] + w;
    else if (val[to] + w >= m2) m2 = val[to] + w;
  }

  for (auto [to,w]: v[x]){
    if (to == par) continue;
    int m = m1;
    if (val[to] + w == m){
      m = m2;
    }

    if (m >= 0){
      upd(m, -1);
      upd(m + w,+1);
    }
    upd(val[to] + w,-1);
    upd(val[to], +1);

    DFS(to,x,m + w);

    upd(val[to], -1);
    upd(val[to] + w,+1);

    if (m >= 0){
      upd(m + w,-1);
      upd(m, +1);
    }
  }
}

void printans(){
  for (int i=1;i<=n;i++)
    cout<<ans[i]<<endl;
}
signed main() {
  ios_base::sync_with_stdio(0),cin.tie(NULL),cout.tie(NULL);
  
  cin>>n>>k;
  for (int i = 1; i < n; i++){
    int a,b,c;
    cin>>a>>b>>c;
    v[a].pb({b,c});
    v[b].pb({a,c});
  }
  
  dfs(1,1);
  upd(val[1],+1);
  upd(0,+1);
  DFS(1,1,0);
  printans();
}

Compilation message (stderr)

Main.cpp: In function 'void upd(long long int, long long int)':
Main.cpp:30:20: 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]
   30 |   while (s1.size() < k && s2.size()){
      |          ~~~~~~~~~~^~~
Main.cpp:37:20: 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 (s1.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