Submission #1361951

#TimeUsernameProblemLanguageResultExecution timeMemory
1361951mariaclaraPetrol stations (CEOI24_stations)C++20
0 / 100
135 ms18088 KiB
#include<bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef pair<int,int> pii;
typedef tuple<int,int,int> trio;
typedef vector<int> vi;
const int MAXN = 7e4+5;
#define all(x) x.begin(), x.end()
#define sz(x) (int)x.size()
#define mk make_pair 
#define pb push_back
#define fr first 
#define sc second

int K;
ll ans[MAXN];
vector<vector<pii>> edges;

int T, allowed[MAXN], euler[MAXN], tin[MAXN], tout[MAXN], pp[MAXN];
ll niv[MAXN];

void construct_euler(int x) {
    euler[T] = x;
    tin[x] = T++;

    for(auto [viz, p] : edges[x]) {
        if(!allowed[viz] or viz == pp[x]) continue;

        niv[viz] = niv[x] + p;
        pp[viz] = x;
        construct_euler(viz);
    }

    tout[x] = T;
}

int anc[MAXN];
vector<pair<ll,int>> att;

void calc_niv(int x) {
    att.pb({niv[x], x});

    if(niv[x] > K) {
        int id = lower_bound(all(att), mk(niv[x] - K, 0)) - att.begin();
        anc[x] = att[id-1].sc;
    }

    for(auto [viz, p] : edges[x]) {
        if(!allowed[viz] or viz == pp[x]) continue;

        niv[viz] = niv[x] + p;
        pp[viz] = x;
        calc_niv(viz);
    }

    att.pop_back();
}

ll rsp[MAXN], bit[MAXN];

void update(int x, ll val) {
    for(x++; x <= T; x += x&-x)
        bit[x] += val;
}

int query(int x) {
    int sum = 0;
    x++;

    while(x > 0) {
        sum += bit[x];
        x -= x&-x;
    }

    return sum;
}

void att_ans(vi a, vi b, int c) {
    for(auto it : a) allowed[it] = 1;
    T = niv[c] = 0;
    pp[c] = -1;

    construct_euler(c);
    
    for(auto it : a) allowed[it] = 0;

    vector<tuple<ll, int, int>> sum;
    vector<pair<ll,int>> ord;

    for(auto x : a) {
        if(x == c) continue;
        sum.pb({K+niv[pp[x]], x, -1});
        sum.pb({K+niv[x], x, 1});
        ord.pb({niv[x], x});
    }

    sort(all(sum));
    sort(all(ord));

    for(int i = 0, j = 0; i < sz(sum); i++) {
        auto [lev, x, fl] = sum[i];

        while(j < sz(ord) and ord[j].fr <= lev) 
            update(tin[ord[j].sc], rsp[ord[j].sc]+1), j++;
        j--;

        rsp[x] += (query(tout[x]) - query(tin[x]-1)) * fl;
    }

    for(auto it : a) rsp[it]++;
    for(auto it : b) allowed[it] = 1, anc[it] = -1;
    niv[c] = 0;
    pp[c] = -1;

    calc_niv(c);

    sum.clear();

    for(auto x : b) {
        if(x == c) continue;
        sum.pb({K - niv[x], x, -1});
        sum.pb({K - niv[pp[x]], x, 1});
    }

    sort(all(sum));

    fill(bit, bit+T+1, 0);

    for(int i = 0, j = 0; i < sz(sum); i++) {
        auto [lev, x, fl] = sum[i];

        while(j < sz(ord) and ord[j].fr <= lev) 
            update(tin[ord[j].sc], rsp[ord[j].sc]), j++;
        j = max(j-1, 0);

        rsp[x] += query(T-1) * fl;
    }

    ord.clear();
    for(auto it : b) ord.pb({niv[it], it});
    sort(all(ord));

    for(auto [lev, x] : ord) 
        if(anc[x] != -1) rsp[x] += rsp[anc[x]];

    for(auto x : b)
        ans[x] += rsp[x];

    for(auto it : b) allowed[it] = 0;
    fill(bit, bit+T+1, 0);
    for(int it : a) rsp[it] = 0;
    for(int it : b) rsp[it] = 0;
}

int blocked[MAXN], sub[MAXN];

void calc_sub(int x, int pai) {
    sub[x] = 1;

    for(auto [viz, p] : edges[x]) {
        if(viz == pai or blocked[viz]) continue;
        calc_sub(viz, x);
        sub[x] += sub[viz];
    }
}

int calc_centroid(int x, int tot, int pai) {
    for(auto [viz, p] : edges[x]) {
        if(viz == pai or blocked[viz]) continue;
        if(sub[x]*2 > tot) return calc_centroid(viz, tot, x);
    }
    return x;
}

vi A, B;

void addA(int x) {
    A.pb(x);

    for(auto [viz, p] : edges[x]) {
        if(blocked[viz] or sub[viz] > sub[x]) continue;
        addA(viz);
    }
}

void addB(int x) {
    B.pb(x);

    for(auto [viz, p] : edges[x]) {
        if(blocked[viz] or sub[viz] > sub[x]) continue;
        addB(viz);
    }
}

void decomp(vi v) {
    if(sz(v) <= 2) return;

    for(auto it : v) blocked[it] = 0, sub[it] = 0;

    calc_sub(v[0], -1);
    int c = calc_centroid(v[0], sz(v), -1);

    for(auto it : v) sub[it] = 0;
    calc_sub(c, -1);

    A.clear(); B.clear();
    int Sa = 0, Sb = 0;
    A.pb(0);
    B.pb(0);

    for(auto [viz, p] : edges[c]) {
        if(blocked[viz]) continue;

        if(abs(Sa+sub[viz] - Sb) <= abs(Sb+sub[viz] - Sa)) 
            addA(viz), Sa += sub[viz];
        else addB(viz), Sb += sub[viz];
    }

    for(auto it : v) blocked[it] = 1, sub[it] = 0;

    att_ans(A, B, c);
    att_ans(B, A, c);

    decomp(A);
    decomp(B);
}

int32_t main() {

    int n;
    cin >> n >> K;

    vi v = {n-1};
    edges.resize(n);

    for(int i = 1, a, b, c; i < n; i++) {
        cin >> a >> b >> c;
        edges[a].pb({b, c});
        edges[b].pb({a, c});
        v.pb(i-1);
    }

    decomp(v);

    for(int i = 0; i < n; i++)
        cout << ans[i] << "\n";
}
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...
#Result Execution timeMemoryGrader output
Fetching results...