oj.uz

Submission #1048066

# Submission time Handle Problem Language Result Execution time Memory
1048066 2024-08-07T21:03:30 Z PurpleCrayon Petrol stations (CEOI24_stations) C++17
26 / 100
 3500 ms 14596 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) { // TODO
        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 sub_all(ll x) { // TODO
    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

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

Subtask #0
0.0 / 0.0

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

Subtask #1
18.0 / 18.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Correct 0 ms 3164 KB Output is correct
3 Correct 3 ms 3164 KB Output is correct
4 Correct 3 ms 3404 KB Output is correct
5 Correct 3 ms 3420 KB Output is correct
6 Correct 2 ms 3420 KB Output is correct
7 Correct 3 ms 3420 KB Output is correct
8 Correct 1 ms 3164 KB Output is correct
9 Correct 3 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 3408 KB Output is correct
14 Correct 3 ms 3164 KB Output is correct

Subtask #2
8.0 / 8.0

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

Subtask #3
0.0 / 10.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Correct 0 ms 3164 KB Output is correct
3 Correct 1 ms 3160 KB Output is correct
4 Correct 2947 ms 13684 KB Output is correct
5 Correct 976 ms 14496 KB Output is correct
6 Execution timed out 3587 ms 14596 KB Time limit exceeded
7 Halted 0 ms 0 KB -

Subtask #4
0.0 / 12.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Correct 0 ms 3164 KB Output is correct
3 Correct 431 ms 7980 KB Output is correct
4 Correct 911 ms 12964 KB Output is correct
5 Execution timed out 3527 ms 13236 KB Time limit exceeded
6 Halted 0 ms 0 KB -

Subtask #5
0.0 / 17.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Correct 0 ms 3164 KB Output is correct
3 Correct 431 ms 7980 KB Output is correct
4 Correct 911 ms 12964 KB Output is correct
5 Execution timed out 3527 ms 13236 KB Time limit exceeded
6 Halted 0 ms 0 KB -

Subtask #6
0.0 / 35.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Correct 0 ms 3164 KB Output is correct
3 Correct 3 ms 3164 KB Output is correct
4 Correct 3 ms 3404 KB Output is correct
5 Correct 3 ms 3420 KB Output is correct
6 Correct 2 ms 3420 KB Output is correct
7 Correct 3 ms 3420 KB Output is correct
8 Correct 1 ms 3164 KB Output is correct
9 Correct 3 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 3408 KB Output is correct
14 Correct 3 ms 3164 KB Output is correct
15 Correct 1 ms 3160 KB Output is correct
16 Correct 2947 ms 13684 KB Output is correct
17 Correct 976 ms 14496 KB Output is correct
18 Execution timed out 3587 ms 14596 KB Time limit exceeded
19 Halted 0 ms 0 KB -