Submission #711085

#	Time	Username	Problem	Language	Result	Execution time	Memory
711085		TimDee	Paths (RMI21_paths)	C++17	19 / 100	368 ms	71080 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<int,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);
int ans[N];

void dfs(int u, int p) {
    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});
        }
    }
    if (adj[u].size()==1) mx[u]={0,u};
}
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;
        auto 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 int sz=1<<30;

struct sgt {
    sgt *L, *R;
    int c,s;
    sgt() {
        L=R=NULL;
        c=s=0;
    }
    void add(int l, int r, int x) {
        c++, s+=x;
        if (r-l==1) return;
        int m=(l+r)>>1;
        if (x<m) {
            if (!L) L=new sgt();
            L->add(l,m,x);
        } else {
            if (!R) R=new sgt();
            R->add(m,r,x);
        }
    }
    void add(int x) {
        add(0,sz,x);
    }
    void del(int l, int r, int x) {
        c--, s-=x;
        if (r-l==1) return;
        int m=(l+r)>>1;
        if (x<m) {
            assert(L!=NULL);
            L->del(l,m,x);
        } else {
            assert(R!=NULL);
            R->del(m,r,x);
        }
    }
    void del(int x) {
        del(0,sz,x);
    }
    int query(int l, int r, int k) {
        if (!k) return 0;
        if (c==k) return s;
        if (r-l==1) return k*l;
        assert(k<=c);
        int ret=0;
        int m=(l+r)>>1;
        if (R) {
            if (R->c < k) {
                k-=R->c;
                ret+=R->s;
            } else return R->query(m,r,k);
        }
        assert(L);
        return ret+L->query(l,m,k);
    }
    int query(int k) {
        return query(0,sz,k);
    }
};

sgt st;
int n,k;
int Z[N];
void dfs3(int u, int p) {
    ans[u]=st.query(k);
    for(auto&e:adj[u]) {
        int v=e.f, w=e.s;
        if (v==p) continue;
        st.del(Z[out[v].s]);
        Z[out[v].s]+=w;
        st.add(Z[out[v].s]);
        st.del(Z[mx[v].s]);
        Z[mx[v].s]-=w;
        st.add(Z[mx[v].s]);
        dfs3(v,u);
        st.del(Z[out[v].s]);
        Z[out[v].s]-=w;
        st.add(Z[out[v].s]);
        st.del(Z[mx[v].s]);
        Z[mx[v].s]+=w;
        st.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) {
        int 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;
        st.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 solve()':
Main.cpp:169:17: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'long long int' [-Wsign-compare]
  169 |     if (z.size()<=k) {
      |         ~~~~~~~~^~~

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