Submission #539985

#	Time	Username	Problem	Language	Result	Execution time	Memory
539985		doowey	Paths (RMI21_paths)	C++14	100 / 100	504 ms	38736 KiB

Main

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;

#define fi first
#define se second
#define mp make_pair
#define fastIO ios::sync_with_stdio(false);cin.tie(0);cout.tie(0);

const int N = (int)1e5 + 10;
vector<pll> T[N];

int k;
struct magic_set{
    set<pll> ms;
    int cnt = 0;
    pll V = mp(-1,-1);
    ll sum = 0;
    // sum of all elements >= V
    void relocate(){
        auto it = ms.lower_bound(V);
        while(cnt > k){
            sum -= it->fi;
            cnt -- ;
            it ++ ;
            V = *it;
        }
        while(cnt < k){
            if(it == ms.begin()) break;
            it -- ;
            sum += it->fi;
            cnt ++ ;
            V = *it;
        }
    }
    void ins(pll X){
        ms.insert(X);
        if(X >= V){
            sum += X.fi;
            cnt++;
        }
        relocate();
    }
    void del(pll X){
        if(X >= V){
            sum -= X.fi;
            cnt -- ;
        }
        ms.erase(X);
        relocate();
    }
};

magic_set Z[N];
pll A[N], B[N];

magic_set G;

void dfs0(int u, int p){
    Z[u].ins(mp(0ll, u));
    for(auto x : T[u]){
        if(x.fi == p) continue;
        dfs0(x.fi, u);

        auto it = Z[x.fi].ms.end();
        it -- ;
        A[x.fi] = *it;
        B[x.fi] = mp(it->fi + x.se, it->se);
        Z[x.fi].del(A[x.fi]);
        Z[x.fi].ins(B[x.fi]);
        if(Z[u].ms.size() < Z[x.fi].ms.size()){
            swap(Z[u], Z[x.fi]);
        }
        for(auto x : Z[x.fi].ms){
            Z[u].ins(x);
        }
        Z[x.fi].ms.clear();
        Z[x.fi].sum = 0;
        Z[x.fi].cnt = 0;
    }
}

ll sol[N];

void reroot(int u, int v, pll go){
    sol[u] = G.sum;
    pll m1 = mp(0, u);
    pll m2 = mp(0, u);
    int a1 = u;
    int a2 = u;
    pll nx;
    for(auto x : T[u]){
        if(x.fi == v) continue;
        if(B[x.fi] > m1){
            m2 = m1;
            a2 = a1;
            m1 = B[x.fi];
            a1 = x.fi;
        }
        else if(B[x.fi] > m2){
            m2 = B[x.fi];
            a2 = x.fi;
        }
    }
    pll d1, i1;
    pll d2, i2;
    for(auto x : T[u]){
        if(x.fi == v) continue;
        d1 = B[x.fi];
        i1 = A[x.fi];
        nx = go;
        if(x.fi != a1){
            nx = max(nx, m1);
        }
        else{
            nx = max(nx, m2);
        }
        d2 = nx;
        i2 = mp(nx.fi + x.se, nx.se);
        G.del(d1);
        G.del(d2);
        G.ins(i1);
        G.ins(i2);
        reroot(x.fi, u, i2);
        G.del(i1);
        G.del(i2);
        G.ins(d1);
        G.ins(d2);
    }
}

int main(){
    fastIO;
    //freopen("in.txt","r",stdin);
    int n;
    cin >> n >> k;
    int u, v, w;
    for(int i = 1; i < n; i ++ ){
        cin >> u >> v >> w;
        T[u].push_back(mp(v, w));
        T[v].push_back(mp(u, w));
    }
    dfs0(1, -1);
    G = Z[1];
    reroot(1, -1, mp(-1ll, -1ll));
    for(int i = 1; i <= n; i ++ ){
        cout << sol[i] << "\n";
    }
    return 0;
}

Compilation message (stderr)

Main.cpp: In function 'void reroot(int, int, pll)':
Main.cpp:94:9: warning: variable 'a2' set but not used [-Wunused-but-set-variable]
   94 |     int a2 = u;
      |         ^~

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