Submission #1106227

#	Time	Username	Problem	Language	Result	Execution time	Memory
1106227		Semen07	Petrol stations (CEOI24_stations)	C++17	100 / 100	492 ms	18748 KiB

Main

#include <bits/stdc++.h>
#define all(a) a.begin(), a.end()
#define int long long
    
using namespace std;
    
template<typename T>
istream &operator>>(istream &in, vector<T> &v) {
    for (auto &u: v) in >> u;
    return in;
}
    
int n, k;
vector<vector<pair<int, int>>> g;
vector<long long> res, w;
vector<int> used, sz, st_to, cnt_to, cnt_from;
vector<pair<long long, int>> srt, st_from;
    
void calc_sz(int v, int p, long long wt) {
    sz[v] = 1;
    w[v] = wt;
    cnt_to[v] = cnt_from[v] = 0;
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        calc_sz(u, v, wt + c);
        sz[v] += sz[u];
    }
}
    
int find_centroid(int v, int p, int full) {
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        if (2 * sz[u] > full) return find_centroid(u, v, full);
    }
    return v;
}
    
void to_centroid(int v, int p, int paths, int mode) {
    st_to.push_back(p);
    cnt_to[v] = 0;
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        to_centroid(u, v, paths, mode);
    }
    res[v] += cnt_to[v] * paths * mode;
    int l = -1, r = (int)st_to.size();
    while (r - l > 1) {
        int mid = (l + r) / 2;
        if (w[v] - w[st_to[mid]] > k) l = mid;
        else r = mid;
    }
    if (r == 0) srt.emplace_back(w[v], cnt_to[v] + 1);
    else cnt_to[st_to[r]] += cnt_to[v] + 1;
    st_to.pop_back();
}
    
void from_centroid_up_to_k(int v, int p, int mode) {
    int left = lower_bound(all(srt), make_pair(k - w[v] + 1, 0ll)) - srt.begin();
    int right = lower_bound(all(srt), make_pair(k - w[p] + 1, 0ll)) - srt.begin();
    int sm = (right > 0) ? srt[right - 1].second: 0;
    if (left > 0) sm -= srt[left - 1].second;
    cnt_from[v] += sm * mode;
    res[p] += sm * sz[v] * mode;
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        from_centroid_up_to_k(u, v, mode);
    }
}

void from_centroid_more_k(int v, int p) {
    int left, right;
    int l = -1, r = (int)st_from.size();
    while (r - l > 1) {
        int mid = (l + r) / 2;
        if (w[p] - st_from[mid].first > k) l = mid;
        else r = mid;
    }
    left = r;
    l = -1, r = (int)st_from.size();
    while (r - l > 1) {
        int mid = (l + r) / 2;
        if (w[v] - st_from[mid].first > k) l = mid;
        else r = mid;
    }
    right = l;
    if (right >= left && right > -1 && left < (int)st_from.size()) {
        int sm = (right >= 0) ? st_from[right].second: 0;
        if (left > 0) sm -= st_from[left - 1].second;
        cnt_from[v] += sm;
        res[p] += sm * sz[v];
    }
    st_from.emplace_back(w[p], (st_from.empty() ? 0: st_from.back().second) + cnt_from[v]);
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        from_centroid_more_k(u, v);
    }
    st_from.pop_back();
}
    
void build_centroid(int v) {
    calc_sz(v, v, 0);
    v = find_centroid(v, v, sz[v]);
    calc_sz(v, v, 0);
    srt.clear();
    srt.emplace_back(0, 1);
    for (auto [u, c]: g[v]) {
        if (!used[u]) to_centroid(u, v, sz[v] - sz[u], 1);
    }
    sort(all(srt));
    for (int i = 1; i < (int)srt.size(); ++i) {
        srt[i].second += srt[i - 1].second;
    }
    for (auto [u, c]: g[v]) {
        if (!used[u]) from_centroid_up_to_k(u, v, 1);
    }
    for (auto [u, c]: g[v]) {
        if (used[u]) continue;
        srt.clear();
        to_centroid(u, v, sz[v] - sz[u], 0);
        sort(all(srt));
        for (int i = 1; i < (int)srt.size(); ++i) {
            srt[i].second += srt[i - 1].second;
        }
        from_centroid_up_to_k(u, v, -1);
    }
    for (auto [u, c]: g[v]) {
        if (!used[u]) from_centroid_more_k(u, v);
    }
    used[v] = 1;
    for (auto [u, c]: g[v]) {
        if (!used[u]) build_centroid(u);
    }
}
    
void solve() {
    cin >> n >> k;
    g.assign(n, {});
    for (int _ = 0; _ < n - 1; ++_) {
        int u, v, c;
        cin >> u >> v >> c;
        g[u].emplace_back(v, c);
        g[v].emplace_back(u, c);
    }
    res.assign(n, 0);
    used.assign(n, 0);
    sz.assign(n, 0);
    w.assign(n, 0);
    cnt_to.assign(n, 0);
    cnt_from.assign(n, 0);
    st_to.clear();
    st_from.clear();
    build_centroid(0);
    for (int i = 0; i < n; ++i) {
        cout << res[i] << '\n';
    }
    cout << '\n';
}
    
signed main() {
    ios_base::sync_with_stdio(false); cin.tie(nullptr);
    int _t = 1;
#ifdef LOCAL
    freopen("io/main.in", "r", stdin);
    freopen("io/main.out", "w", stdout);
    cin >> _t;
#endif
    while (_t--) solve();
}

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