Submission #1224785

#	Time	Username	Problem	Language	Result	Execution time	Memory
1224785		Mighilon	Paths (RMI21_paths)	C++20	100 / 100	341 ms	39976 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
#ifdef DEBUG
#include "../Library/debug.h"
#else
#define dbg(x...)
#endif
 
#define int long long
typedef long long ll;
typedef long double ld;
typedef pair<int, int> pi;
typedef pair<ll, ll> pl;
typedef vector<int> vi;
typedef vector<bool> vb;
typedef vector<ll> vl;
typedef vector<pi> vpi;
typedef vector<pl> vpl; 
 
#define FOR(i, a, b) for (int i = (a); i < (b); ++i)
#define F0R(i, a) for (int i = 0; i < (a); ++i)
#define FORd(i, a, b) for (int i = (b) - 1; i >= (a); --i)
#define F0Rd(i, a) for (int i = (a) - 1; i >= 0; --i)
#define trav(a, x) for (auto& a : x)
#define f first 
#define s second
#define pb push_back
#define sz(x) (int)(x).size()
#define all(x) x.begin(), x.end()
 
const char nl = '\n';
const int INF = 1e9;
// const ll MOD = 998244353;
const int mxN = -1;

struct SG{
  int n;
  using T = pair<int,int>;
  vector<T> t;
  SG(){}
  SG(int _n){
    n=_n;
    if(__builtin_popcount(n)!=1)
      n=1<<(32-__builtin_clz(n));
    t.resize(2*n);
  }
  T merge(T a, T b){
    if(a>b)return a;
    return b;
  }
  void add(int x, int l, int r, int i, ll val){
    if(l>i||r<i)return;
    if(l==r&&l==i){
      t[x].f+=val;
      t[x].s=i;
      return;
    }
    int m=(l+r)/2;
    add(2*x,l,m,i,val);
    add(2*x+1,m+1,r,i,val);
    t[x]=merge(t[2*x],t[2*x+1]);
  }
  void add(int i, ll val){
    add(1,0,n-1,i,val);
  }
  T query(int x, int l, int r, int ql, int qr){
    if(l>qr||r<ql){
      return {0,0};
    }
    if(l>=ql&&r<=qr){
      return t[x];
    }
    int m=(l+r)/2;
    return merge(query(2*x,l,m,ql,qr),query(2*x+1,m+1,r,ql,qr));
  }
  T query(int l, int r){
    if(l>r)return{-1,-1};
    return query(1,0,n-1,l,r);
  }
  T query_out(int l, int r){
    return merge(query(0,l-1),query(r+1,n-1));
  }
}sg;

vector<vi> adj;
vector<set<pair<ll,ll>>> st;
vl e,c;
int n,k;
vpi in;
vi id;
int id_cnt=0;

void initdfs(int v, int p){
  trav(_i,adj[v]){
    int u=e[_i]^v;
    if(u==p)continue;
    initdfs(u,v);
  }
  if(sz(adj[v])==1){
    id[v]=id_cnt++;
  }
}


void dfs(int v,int p){
  in[v]={1e9,-1e9};
  trav(_i,adj[v]){
    int u=e[_i]^v;
    if(u==p)continue;
    dfs(u,v);
    auto it=prev(st[u].end());
    pl x=*it; x.f+=c[_i];
    st[u].erase(it);
    st[u].insert(x);
    if(sz(st[v])<sz(st[u])) swap(st[v],st[u]);
    trav(i,st[u])
      st[v].insert(i);
    in[v].f=min(in[v].f,in[u].f);
    in[v].s=max(in[v].s,in[u].s);
  }
  if(sz(adj[v])==1){
    st[v].insert({0,id[v]});
    in[v]={id[v],id[v]};
  }
}


vl ans;
set<pl> sl,sr;

void balance(ll& sum){
  while(sz(sr)<k&&!sl.empty()){
    auto it=prev(sl.end());
    sr.insert(*it);
    sum+=it->f;
    sl.erase(it);
  }
  while(sz(sr)>k){
    auto it=sr.begin();
    sl.insert(*it);
    sum-=it->f;
    sr.erase(it);
  }
  if(sl.empty())return;
  auto it1=prev(sl.end());
  auto it2=sr.begin();;
  while(*it1>*it2){
    pl _it1=*it1,_it2=*it2;
    sl.erase(it1);
    sr.erase(it2);
    sl.insert(_it2);
    sr.insert(_it1);
    it1=prev(sl.end());
    it2=sr.begin();
    sum-=_it2.f;
    sum+=_it1.f;
  }
}
void erase(ll& sum,pl a){
  auto it = sl.find(a);
  if(it!=sl.end())
    sl.erase(it);
  it = sr.find(a);
  if(it!=sr.end()){
    sum-=it->f;
    sr.erase(it);
  }
}

void dfs2(int v, int p,ll sum){
  dbg(v,sum)
  dbg(sr,sl)
  ans[v]=sum;
  trav(_i,adj[v]){
    int u=e[_i]^v;
    if(u==p)continue;

    dbg(sg.t)
    dbg(v,u,sum);
    //in
    pl x=sg.query(in[u].f,in[u].s);
    erase(sum, x);
    dbg(sum,x.f-c[_i],x,in[u]);
    sr.insert({x.f-c[_i],x.s});
    sum+=x.f-c[_i];
    dbg(sum);
    sg.add(x.s,-c[_i]);
    //out
    pl y=sg.query_out(in[u].f,in[u].s);
    erase(sum, y);
    dbg(sum, y)
    sr.insert({y.f+c[_i],y.s});
    sum+=y.f+c[_i];
    dbg(sum,y.f+c[_i]);
    sg.add(y.s,c[_i]);

    balance(sum);
    dbg(sum);

    dfs2(u,v,sum);


    //in
    erase(sum, {x.f-c[_i],x.s});
    sr.insert(x);
    sum+=x.f;
    sg.add(x.s,c[_i]);
    //out
    erase(sum, {y.f+c[_i],y.s});
    sr.insert(y);
    sum+=y.f;
    sg.add(y.s,-c[_i]);
    balance(sum);
    dbg("end",v,u, sum);
  }
}

 
void solve(){
  cin>>n>>k;
  in.resize(n);
  e.resize(n-1);
  c.resize(n-1);
  adj.resize(n);
  ans.resize(n);
  st.resize(n);
  id.resize(n);
  F0R(i,n-1){
    int a,b,co;
    cin>>a>>b>>co;--a,--b;
    e[i]=a^b;
    c[i]=co;
    adj[a].pb(i);
    adj[b].pb(i);
  }
  int s=0;
  F0R(i,n){
    if(sz(adj[i])>1){
      s=i;
      break;
    }
  }
  initdfs(s,-1);
  dfs(s,-1);
  ll sum=0;
  sg=SG(id_cnt);
  dbg(st[s])
  trav(i,st[s]){
    sg.add(i.s, i.f);
  }
  sl=st[s];
  balance(sum);
  dfs2(s,-1,sum);
  trav(i,ans){
    cout<<i<<nl;
  }
}
int32_t main(){
    ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
 
    int TC = 1;
    // cin >> TC;
    while(TC--){
        solve();
    }
    return 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