Submission #1048063

#	Time	Username	Problem	Language	Result	Execution time	Memory
1048063		PurpleCrayon	Petrol stations (CEOI24_stations)	C++17	18 / 100	3594 ms	14352 KiB

Main

#include <bits/stdc++.h>
using namespace std;
 
#define sz(v) int(v.size())
#define ar array
typedef long long ll;
const int N = 7e4+10, MOD = 1e9+7;
const ll INF = 1e18+10;

int n, K;
ar<int, 3> edges[N];
vector<ar<int, 3>> adj[N];
int sub[N];
ll ans[N];

void make_tree(vector<int>& eids) {
    for (int e : eids) {
        auto [a, b, w] = edges[e];
        adj[a].clear(), adj[b].clear();
    }

    for (int e : eids) {
        auto [a, b, w] = edges[e];
        adj[a].push_back({b, w, e});
        adj[b].push_back({a, w, e});
    }
}

int dfs_sub(int c, int p) {
    sub[c] = 1;
    for (auto [nxt, w, eid] : adj[c]) if (nxt != p) {
        sub[c] += dfs_sub(nxt, c);
    }
    return sub[c];
}

int dfs_cent(int c, int p, int k) {
    for (auto [nxt, w, eid] : adj[c]) if (nxt != p && sub[nxt] > k / 2) {
        return dfs_cent(nxt, c, k);
    }
    return c;
}

void add_branch(int c, int p, vector<int>& es) {
    for (auto [nxt, w, eid] : adj[c]) if (nxt != p) {
        es.push_back(eid);
        add_branch(nxt, c, es);
    }
}

pair<vector<int>, vector<int>> split_cent(int c) {
    int k = sub[c];

    vector<int> ae, be;
    int sum = 0;
    bool done = 0;
    for (auto [nxt, w, e] : adj[c]) {
        if (done) {
            be.push_back(e);
            continue;
        }

        if (k / 3 <= sub[nxt] && sub[nxt] <= 2 * k / 3) {
            be = ae;
            ae = {e};
            done = 1;
            continue;
        }

        if (sum < k / 3) {
            sum += sub[nxt];
            ae.push_back(e);
        } else {
            done = 1;
            be.push_back(e);
        }
    }

    int sa = sz(ae);
    for (int i = 0; i < sa; i++) {
        int nxt = edges[ae[i]][0] ^ edges[ae[i]][1] ^ c;
        add_branch(nxt, c, ae);
    }

    int sb = sz(be);
    for (int i = 0; i < sb; i++) {
        int nxt = edges[be[i]][0] ^ edges[be[i]][1] ^ c;
        add_branch(nxt, c, be);
    }

    return make_pair(ae, be);
}

ll stock[N];
vector<pair<ll, ll>> cur_extra;
vector<ll> stk_w;
vector<int> stk_anc;
void dfs_stock(int c, int p, int he_sz) {
    stock[c] = 0;
    stk_anc.push_back(c);
    for (auto [nxt, w, eid] : adj[c]) if (nxt != p) {
        stk_w.push_back(w);
        dfs_stock(nxt, c, he_sz);
        stk_w.pop_back();
    }

    if (p == -1) return;

    stock[c]++;
    // cerr << "add: " << c << ' ' << stock[c] << ' ' << stock[c] * he_sz << endl;
    ans[c] += stock[c] * he_sz;
    ll sum = K;
    int cp = 0;
    while (sum >= 0) {
        sum -= stk_w[sz(stk_w) - 1 - cp];
        cp++;
    }

    // cp = 1 means use stk_anc.back()
    assert(cp >= 2);
    if (cp == sz(stk_anc)) {
        cur_extra.emplace_back(sum + stk_w[0], stock[c]);
    } else {
        stock[stk_anc[sz(stk_anc) - cp]] += stock[c];
    }

    stk_anc.pop_back();
}

void sift_up(int c, int he_sz) {
    cur_extra.clear();
    stk_w = {INF};
    stk_anc.clear();
    dfs_stock(c, -1, he_sz);
}

ll sub_all(ll x) {
    ll sum = 0;
    for (auto& [he, cnt] : cur_extra) {
        if (he >= 0 && he - x < 0) {
            sum += cnt;
        }

        he -= x;
    }
    return sum;
}

void ins(ll x, ll cnt) {
    cur_extra.emplace_back(x, cnt);
}

void er(ll x, ll cnt) {
    auto it = find(cur_extra.begin(), cur_extra.end(), make_pair(x, cnt));
    assert(it != cur_extra.end());
    cur_extra.erase(it);
}

void add_all(ll x) {
    for (auto& [he, cnt] : cur_extra)
        he += x;
}

void dfs_down(int c, int p) {
    for (auto [nxt, w, eid] : adj[c]) if (nxt != p) {
        ll cnt = sub_all(w); // returns the number that are now < 0
        // cerr << "dfs_down: " << c << ' ' << nxt << ' ' << w << ' ' << cnt << endl;
        ans[c] += cnt * sub[nxt];
        if (cnt) ins(K - w, cnt);
        dfs_down(nxt, c);
        if (cnt) er(K - w, cnt);
        add_all(w);
    }
}

void sift_down(int c) { // use cur_extra
    sort(cur_extra.begin(), cur_extra.end());
    dfs_sub(c, -1);
    dfs_down(c, -1);
}

void run_cent(int c, vector<int> ae, vector<int> be) {
    // cerr << "sifting up: " << ae[0] << ' ' << "down: " << be[0] << endl;
    make_tree(ae);
    sift_up(c, sz(be)); // exclude c
    // for (auto [x, cnt] : cur_extra) cerr << "cur_extra: " << x << ' ' << cnt << endl;
    make_tree(be);
    sift_down(c);
}

void build(vector<int>& eids) {
    // cerr << "eids: ";
    // for (int e : eids) cerr << e << ' ';
    // cerr << endl;
    if (sz(eids) <= 1) return; // doesn't count (a, a) or single edges (a, b)
    make_tree(eids);
    int k = dfs_sub(edges[eids[0]][0], -1);
    int c = dfs_cent(edges[eids[0]][0], -1, k);
    // cerr << "centroid: " << c << endl << endl;

    dfs_sub(c, -1);
    auto [a, b] = split_cent(c);

    run_cent(c, a, b);
    // cerr << endl;
    // for (int i = 0; i < n; i++) cerr << ans[i] << ' ';
    // cerr << endl;
    run_cent(c, b, a);
    build(a), build(b);
}

void solve() {
    cin >> n >> K;
    for (int i = 0; i < n-1; i++) {
        int a, b, w; cin >> a >> b >> w;
        edges[i] = {a, b, w};
    }

    vector<int> eids(n-1);
    iota(eids.begin(), eids.end(), 0);
    build(eids);

    assert(sz(eids) == n-1);
    make_tree(eids);
    for (int i = 0; i < n; i++) {
        ans[i]++;
        for (auto [nxt, w, eid] : adj[i]) {
            ans[i]++;
        }
    }

    for (int i = 0; i < n; i++) {
        ans[i] -= n;
        cout << ans[i] << '\n';
    }
}
 
int main() {
    ios::sync_with_stdio(false); cin.tie(0);
    int T = 1;
    // cin >> T;
    while (T--) solve();
}

Compilation message (stderr)

Main.cpp: In function 'void solve()':
Main.cpp:227:19: warning: structured binding declaration set but not used [-Wunused-but-set-variable]
  227 |         for (auto [nxt, w, eid] : adj[i]) {
      |                   ^~~~~~~~~~~~~

• Subtask 0: 0 / 0
• Subtask 1: 18 / 18
• Subtask 2: 0 / 8
• Subtask 3: 0 / 10
• Subtask 4: 0 / 12
• Subtask 5: 0 / 17
• Subtask 6: 0 / 35