oj.uz

Submission #1048063

# Submission time Handle Problem Language Result Execution time Memory
1048063 2024-08-07T20:58:35 Z PurpleCrayon Petrol stations (CEOI24_stations) C++17
18 / 100
 3500 ms 14352 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;
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

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 / 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 2 ms 3304 KB Output is correct
4 Correct 3 ms 3164 KB Output is correct
5 Correct 4 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 2 ms 3420 KB Output is correct
10 Correct 3 ms 3420 KB Output is correct
11 Correct 2 ms 3420 KB Output is correct
12 Correct 3 ms 3420 KB Output is correct
13 Correct 3 ms 3420 KB Output is correct
14 Correct 3 ms 3420 KB Output is correct

Subtask #2
0.0 / 8.0

# Verdict Execution time Memory Grader output
1 Correct 1 ms 3164 KB Output is correct
2 Execution timed out 3594 ms 14352 KB Time limit exceeded
3 Halted 0 ms 0 KB -

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 3164 KB Output is correct
4 Execution timed out 3594 ms 14352 KB Time limit exceeded
5 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 485 ms 8716 KB Output is correct
4 Correct 1108 ms 13596 KB Output is correct
5 Execution timed out 3563 ms 14144 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 485 ms 8716 KB Output is correct
4 Correct 1108 ms 13596 KB Output is correct
5 Execution timed out 3563 ms 14144 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 2 ms 3304 KB Output is correct
4 Correct 3 ms 3164 KB Output is correct
5 Correct 4 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 2 ms 3420 KB Output is correct
10 Correct 3 ms 3420 KB Output is correct
11 Correct 2 ms 3420 KB Output is correct
12 Correct 3 ms 3420 KB Output is correct
13 Correct 3 ms 3420 KB Output is correct
14 Correct 3 ms 3420 KB Output is correct
15 Correct 1 ms 3164 KB Output is correct
16 Execution timed out 3594 ms 14352 KB Time limit exceeded
17 Halted 0 ms 0 KB -