oj.uz

Submission #1048068

# Submission time Handle Problem Language Result Execution time Memory
1048068 2024-08-07T21:08:38 Z PurpleCrayon Petrol stations (CEOI24_stations) C++17
100 / 100
 1468 ms 12732 KB 
#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, ps_stk;
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);
        ps_stk.push_back(ps_stk.back() + w);
        dfs_stock(nxt, c, he_sz);
        stk_w.pop_back();
        ps_stk.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++;
    }
    */

    // smallest number s.t. sum > K
    int lo = -1, hi = sz(stk_w) + 1, mid = (lo + hi) / 2;
    while (lo < mid && mid < hi) {
        if (ps_stk.back() - ps_stk[sz(ps_stk) - mid - 1] > K) hi = mid;
        else lo = mid;
        mid = (lo + hi) / 2;
    }

    int cp = hi;
    ll sum = K - (ps_stk.back() - ps_stk[sz(ps_stk) - hi - 1]);

    // 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};
    ps_stk = {0, INF};
    stk_anc.clear();
    dfs_stock(c, -1, he_sz);
}

ll delta = 0;

ll sub_all(ll x) {
    // need the count in the range [0, x-1]
    // [-delta, -delta + x-1]
    
    int L = lower_bound(cur_extra.begin(), cur_extra.end(), make_pair(-delta, -INF)) - cur_extra.begin();
    int R = upper_bound(cur_extra.begin(), cur_extra.end(), make_pair(-delta + x-1, INF)) - cur_extra.begin() - 1;
    delta -= x;

    ll sum = 0;
    for (int i = L; i <= R; i++)
        sum += cur_extra[i].second;

    /*
    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) {
    x -= delta;
    cur_extra.emplace_back(x, cnt);
}

void er(ll x, ll cnt) {
    x -= delta;
    cur_extra.pop_back();
}

void add_all(ll x) {
    delta += 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

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

Subtask #0
0.0 / 0.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3160 KB Output is correct
2 Correct 1 ms 3164 KB Output is correct

Subtask #1
18.0 / 18.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3160 KB Output is correct
2 Correct 1 ms 3164 KB Output is correct
3 Correct 2 ms 3160 KB Output is correct
4 Correct 2 ms 3164 KB Output is correct
5 Correct 2 ms 3416 KB Output is correct
6 Correct 3 ms 3420 KB Output is correct
7 Correct 2 ms 3420 KB Output is correct
8 Correct 1 ms 3164 KB Output is correct
9 Correct 2 ms 3164 KB Output is correct
10 Correct 2 ms 3164 KB Output is correct
11 Correct 2 ms 3164 KB Output is correct
12 Correct 2 ms 3164 KB Output is correct
13 Correct 2 ms 3164 KB Output is correct
14 Correct 2 ms 3164 KB Output is correct

Subtask #2
8.0 / 8.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Correct 244 ms 12732 KB Output is correct

Subtask #3
10.0 / 10.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3160 KB Output is correct
2 Correct 1 ms 3164 KB Output is correct
3 Correct 1 ms 3164 KB Output is correct
4 Correct 244 ms 12732 KB Output is correct
5 Correct 253 ms 12192 KB Output is correct
6 Correct 253 ms 12728 KB Output is correct

Subtask #4
12.0 / 12.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3160 KB Output is correct
2 Correct 1 ms 3164 KB Output is correct
3 Correct 234 ms 7808 KB Output is correct
4 Correct 259 ms 12092 KB Output is correct
5 Correct 242 ms 12732 KB Output is correct
6 Correct 1 ms 3164 KB Output is correct
7 Correct 231 ms 9428 KB Output is correct
8 Correct 251 ms 9272 KB Output is correct
9 Correct 241 ms 9400 KB Output is correct

Subtask #5
17.0 / 17.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3160 KB Output is correct
2 Correct 1 ms 3164 KB Output is correct
3 Correct 234 ms 7808 KB Output is correct
4 Correct 259 ms 12092 KB Output is correct
5 Correct 242 ms 12732 KB Output is correct
6 Correct 1 ms 3164 KB Output is correct
7 Correct 231 ms 9428 KB Output is correct
8 Correct 251 ms 9272 KB Output is correct
9 Correct 241 ms 9400 KB Output is correct
10 Correct 328 ms 10196 KB Output is correct
11 Correct 215 ms 9440 KB Output is correct
12 Correct 226 ms 9396 KB Output is correct
13 Correct 208 ms 9452 KB Output is correct
14 Correct 229 ms 9500 KB Output is correct
15 Correct 210 ms 9400 KB Output is correct
16 Correct 1137 ms 10840 KB Output is correct

Subtask #6
35.0 / 35.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3160 KB Output is correct
2 Correct 1 ms 3164 KB Output is correct
3 Correct 2 ms 3160 KB Output is correct
4 Correct 2 ms 3164 KB Output is correct
5 Correct 2 ms 3416 KB Output is correct
6 Correct 3 ms 3420 KB Output is correct
7 Correct 2 ms 3420 KB Output is correct
8 Correct 1 ms 3164 KB Output is correct
9 Correct 2 ms 3164 KB Output is correct
10 Correct 2 ms 3164 KB Output is correct
11 Correct 2 ms 3164 KB Output is correct
12 Correct 2 ms 3164 KB Output is correct
13 Correct 2 ms 3164 KB Output is correct
14 Correct 2 ms 3164 KB Output is correct
15 Correct 1 ms 3164 KB Output is correct
16 Correct 244 ms 12732 KB Output is correct
17 Correct 253 ms 12192 KB Output is correct
18 Correct 253 ms 12728 KB Output is correct
19 Correct 234 ms 7808 KB Output is correct
20 Correct 259 ms 12092 KB Output is correct
21 Correct 242 ms 12732 KB Output is correct
22 Correct 1 ms 3164 KB Output is correct
23 Correct 231 ms 9428 KB Output is correct
24 Correct 251 ms 9272 KB Output is correct
25 Correct 241 ms 9400 KB Output is correct
26 Correct 328 ms 10196 KB Output is correct
27 Correct 215 ms 9440 KB Output is correct
28 Correct 226 ms 9396 KB Output is correct
29 Correct 208 ms 9452 KB Output is correct
30 Correct 229 ms 9500 KB Output is correct
31 Correct 210 ms 9400 KB Output is correct
32 Correct 1137 ms 10840 KB Output is correct
33 Correct 248 ms 9640 KB Output is correct
34 Correct 373 ms 10728 KB Output is correct
35 Correct 204 ms 10012 KB Output is correct
36 Correct 231 ms 10036 KB Output is correct
37 Correct 229 ms 10456 KB Output is correct
38 Correct 232 ms 10508 KB Output is correct
39 Correct 237 ms 10376 KB Output is correct
40 Correct 1468 ms 11404 KB Output is correct