제출 #1106062

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1106062		Semen07	Petrol stations (CEOI24_stations)	C++17	55 / 100	3554 ms	40432 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;
}

class SegTree {
    struct node {
        int val = 0;
        node *left = nullptr, *right = nullptr;
        node() = default;
        node(int x): val(x) {}
        void push() {
            if (!left) left = new node();
            if (!right) right = new node();
        }
        ~node() {
            if (left) delete left;
            if (right) delete right;
        }
    };
    int n = 0, shift = 0;
    node *root = nullptr;
    void add(node *p, int l, int r, int i, int x) {
        if (l > i || r <= i) return;
        if (l + 1 == r) {
            p->val += x;
            return;
        }
        int mid = (l + r) / 2; p->push();
        add(p->left, l, mid, i, x);
        add(p->right, mid, r, i, x);
        p->val = p->left->val + p->right->val;
    }
    int get(node *p, int l, int r, int a, int b) {
        if (l >= b || r <= a || p->val == 0) return 0;
        if (a <= l && r <= b) return p->val;
        int mid = (l + r) / 2; p->push();
        return get(p->left, l, mid, a, b) + get(p->right, mid, r, a, b);
    }
 public:
    void assign(int _n) {
        n = _n;
        if (root) delete root;
        root = new node();
        shift = 0;
    }
    SegTree() = default;
    void add(int i, int x) {
        add(root, 0, n, i + shift, x);
    }
    int get(int l, int r) {
        return get(root, 0, n, l + shift, r + shift);
    }
    void make_shift(int x) {
        shift += x;
    }
    ~SegTree() {
        if (root) delete root;
    }
};

constexpr int inf = 1'000'000'000'000;

int n, k;
vector<vector<pair<int, int>>> g;
vector<int> res, used, sz, w, st, cnt;
SegTree d;

void calc_sz(int v, int p, int wt) {
    sz[v] = 1;
    w[v] = wt;
    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 first_dfs(int v, int p, int paths, bool flag) {
    st.push_back(v);
    cnt[v] = 0;
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        first_dfs(u, v, paths, flag);
    }
    int l = -1, r = (int)st.size();
    while (r - l > 1) {
        int mid = (l + r) / 2;
        if (w[v] - w[st[mid]] > k) l = mid;
        else r = mid;
    }
    if (r == 0) d.add(k - w[v], cnt[v] + 1);
    else cnt[st[r]] += cnt[v] + 1;
    if (flag) res[v] += cnt[v] * paths;
    st.pop_back();
}

void second_dfs(int v, int p) {
    for (auto [u, c]: g[v]) {
        if (u == p || used[u]) continue;
        int sm = d.get(0, c);
        res[v] += sm * sz[u];
        d.add(k, sm);
        d.make_shift(c);
        second_dfs(u, v);
        d.make_shift(-c);
        d.add(k, -sm);
    }
}

void build_centroid(int v) {
    calc_sz(v, v, 0);
    int cent = find_centroid(v, v, sz[v]);
    calc_sz(cent, cent, 0);
    used[cent] = 1;

    st = {cent};

    d.assign(inf);
    for (int i = 0; i < (int)g[cent].size(); ++i) {
        auto [u, c] = g[cent][i];
        if (used[u]) continue;
        int sm = d.get(0, c);
        res[cent] += sm * sz[u];
        d.add(k, sm + 1);
        d.make_shift(c);
        second_dfs(u, cent);
        d.make_shift(-c);
        d.add(k, -sm - 1);
        first_dfs(u, cent, sz[cent] - sz[u], true);
    }

    st = {cent};

    d.assign(inf);
    for (int i = (int)g[cent].size() - 1; i >= 0; --i) {
        auto [u, c] = g[cent][i];
        if (used[u]) continue;
        int sm = d.get(0, c);
        res[cent] += sm * sz[u];
        d.add(k, sm);
        d.make_shift(c);
        second_dfs(u, cent);
        d.make_shift(-c);
        d.add(k, -sm);
        first_dfs(u, cent, sz[cent] - sz[u], false);
    }

    for (auto [u, c]: g[cent]) {
        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.assign(n, 0);
    st.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: 0 / 10
• Subtask 4: 12 / 12
• Subtask 5: 17 / 17
• Subtask 6: 0 / 35