Submission #711174

#	Time	Username	Problem	Language	Result	Execution time	Memory
711174		TimDee	Paths (RMI21_paths)	C++17	100 / 100	505 ms	248728 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
#pragma GCC optimize("O3")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx2,popcnt")

using ll = long long;
//#define int long long
//#define double long double
#define forn(i,n) for(int i=0; i<(n); ++i)
#define pb push_back
#define pi pair<ll,int>
#define f first
#define s second 
#define vii(a,n) vector<int> a(n); forn(i,n) cin>>a[i];
#define all(x) x.begin(), x.end()
#define rall(x) x.rbegin(), x.rend()

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
//const int inf = 1e18;
const int mod = 1e9+7;//998244353;

const int N = 1e5;
vector<pi> adj[N];
vector<pi> z;
vector<pi> mx(N,{-1,-1});
vector<pi> out(N);
ll ans[N];

void dfs(int u, int p) {
    if (adj[u].size()==1) mx[u]={0,u};
    for (auto&e:adj[u]) {
        int v=e.f, w=e.s;
        if (v==p) continue;
        dfs(v,u);
        if (mx[v].f+w > mx[u].f) {
            if (mx[u].s!=-1) z.pb(mx[u]);
            mx[u]={mx[v].f+w,mx[v].s};
        } else {
            z.pb({mx[v].f+w,mx[v].s});
        }
    }
}
void dfs2(int u, int p, pi o) {
    out[u]=o;
    vector<pi> pr={{-1,-1}};
    for(auto&e:adj[u]) {
        int v=e.f, w=e.s;
        if (v==p) {
            pr.pb(pr.back());
            continue;
        }
        if (mx[v].f+w>pr.back().f) pr.pb({mx[v].f+w,mx[v].s});
        else pr.pb(pr.back());
    }
    pi sf={-1,-1};
    int sz=adj[u].size();
    for (int i=sz-1; i>=0; --i) {
        auto e=adj[u][i];
        int v=e.f, w=e.s;
        if (v==p) continue;
        pi x=(pr[i].f > sf.f)?pr[i]:sf;
        if (o.f > x.f) x=o;
        x.f+=w;
        dfs2(v,u,x);
        if (mx[v].f+w > sf.f) sf={mx[v].f+w,mx[v].s};
    }
}

const ll sz=1ll<<47;

struct sgt {
    int l,r;
    int c;
    ll s;
    sgt() {
        l=r=-1;
        c=s=0;
    }
};

const int tsz=9.6e6;
sgt t[tsz];
int nxt=1;

void add(int v, ll l, ll r, ll x) {
    t[v].c++, t[v].s+=x;
    if (r-l==1) return;
    ll m=(l+r)>>1;
    int L=t[v].l, R=t[v].r;
    if (x<m) {
        if (L==-1) t[v].l=nxt++;
        add(t[v].l,l,m,x);
    } else {
        if (R==-1) t[v].r=nxt++;
        add(t[v].r,m,r,x);
    }
}
void add(ll x) {
    add(0,0,sz,x);
}
void del(int v, ll l, ll r, ll x) {
    t[v].c--, t[v].s-=x;
    if (r-l==1) return;
    ll m=(l+r)>>1;
    int L=t[v].l, R=t[v].r;
    if (x<m) {
        del(t[v].l,l,m,x);
    } else {
        del(t[v].r,m,r,x);
    }
}
void del(ll x) {
    del(0,0,sz,x);
}

ll query(int v, ll l, ll r, int k) {
    ll c=t[v].c, s=t[v].s;
    if (c==k) return s;
    if (r-l==1) return 1ll*k*l;
    ll ret=0;
    ll m=(l+r)>>1;
    int L=t[v].l, R=t[v].r;
    if (R!=-1) {
        if (t[R].c < k) {
            k-=t[R].c;
            ret+=t[R].s;
        } else return query(R,m,r,k);
    }
    return ret+query(L,l,m,k);
}
ll query(int k) {
    return query(0,0,sz,k);
}

int n,k;
ll Z[N];
void dfs3(int u, int p) {
    ans[u]=query(k);
    for(auto&e:adj[u]) {
        int v=e.f, w=e.s;
        if (v==p) continue;
        del(Z[out[v].s]);
        Z[out[v].s]+=w;
        add(Z[out[v].s]);
        del(Z[mx[v].s]);
        Z[mx[v].s]-=w;
        add(Z[mx[v].s]);
        dfs3(v,u);
        del(Z[out[v].s]);
        Z[out[v].s]-=w;
        add(Z[out[v].s]);
        del(Z[mx[v].s]);
        Z[mx[v].s]+=w;
        add(Z[mx[v].s]);
    }
}

void solve() {

	cin>>n>>k;
    forn(i,n-1) {
        int u,v,w; cin>>u>>v>>w; --u, --v;
        adj[u].pb({v,w});
        adj[v].pb({u,w});
    }
    dfs(0,-1);
    z.pb(mx[0]);


    pi puiu1={0,0}, puiu2={-1,-1};
    pi x = (adj[0].size()==1)?puiu1:puiu2;
    dfs2(0,-1,x);

    if (z.size()<=k) {
        ll s=0;
        for(auto&x:z) s+=x.f;
        forn(i,n) cout<<s<<'\n';
        return;
    }
    for(auto&x:z) {
        Z[x.s]=x.f;
        add(x.f);
    }
    dfs3(0,-1);
    forn(i,n) cout<<ans[i]<<'\n';

}
     
int32_t main() {
	ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    int t=1;
    //cin>>t;
    while (t--) solve();
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void del(int, ll, ll, ll)':
Main.cpp:107:9: warning: unused variable 'L' [-Wunused-variable]
  107 |     int L=t[v].l, R=t[v].r;
      |         ^
Main.cpp:107:19: warning: unused variable 'R' [-Wunused-variable]
  107 |     int L=t[v].l, R=t[v].r;
      |                   ^
Main.cpp: In function 'void solve()':
Main.cpp:176:17: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<long long int, int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  176 |     if (z.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