Submission #1053390

#	Time	Username	Problem	Language	Result	Execution time	Memory
1053390		elazarkoren	Petrol stations (CEOI24_stations)	C++17	0 / 100	71 ms	36572 KiB

Main

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#define x first
#define y second
#define all(v) v.begin(), v.end()
#define chkmin(a, b) a = min(a, b)
#define chkmax(a, b) a = max(a, b)
//#define int ll
using namespace std;
//using namespace __gnu_pbds;
typedef long long ll;
typedef vector<ll> vi;
typedef vector<vi> vvi;
typedef pair<ll, ll> pii;
typedef vector<pii> vii;
typedef vector<bool> vb;
typedef __gnu_pbds::tree<pii, __gnu_pbds::null_type, less<pii>, __gnu_pbds::rb_tree_tag, __gnu_pbds::tree_order_statistics_node_update> ordered_set;

const int MAX_N = 1e5 + 5;

struct Seg {
    vi seg;
    int n;
    Seg() {}
    Seg(int m) {
        seg.resize(2 * n);
    }
    void Update(int i, ll x) {
        seg[i += n] = x;
        for (i >>= 1; i; i >>= 1) {
            seg[i] = seg[i << 1] + seg[i << 1 | 1];
        }
    }
    ll Query(int l, int r) {
        ll ans = 0;
        for (l += n, r += n; l < r; l >>= 1, r >>= 1) {
            if (l & 1) ans += seg[l++];
            if (r & 1) ans += seg[--r];
        }
        return ans;
    }
};

ll ans[MAX_N];
vii tree[MAX_N];
ll sz[MAX_N], real_sz[MAX_N];
bool visited[MAX_N];
int n, k;

int RealDfsSz(int node, int parent) {
    real_sz[node] = 1;
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent && !visited[neighbor]) {
            real_sz[node] += RealDfsSz(neighbor, node);
        }
    }
    return real_sz[node];
}

int DfsSz(int node, int parent) {
    sz[node] = 1;
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent && !visited[neighbor]) {
            sz[node] += DfsSz(neighbor, node);
        }
    }
    return sz[node];
}

int FindCentroid(int node) {
    int s = DfsSz(node, -1);
    int par = node;
    while (true) {
        bool ok = true;
        for (auto [neighbor, w] : tree[node]) {
            if (visited[neighbor] || neighbor == par) continue;
            if (sz[neighbor] > s / 2) {
                par = node;
                node = neighbor;
                ok = false;
                break;
            }
        }
        if (ok) break;
    }
    return node;
}

int tim;
ordered_set se;

int dp[MAX_N];
void CalcDp(int node, int parent, int d, vii &stk) {
    if (d <= k) {
        dp[node] = d;
    } else {
        int begin = 0, end = stk.size(), mid;
        while (begin < end) {
            mid = (begin + end) >> 1;
            if (d - stk[mid].x <= k) end = mid;
            else begin = mid + 1;
        }
        dp[node] = stk[end].y;
    }
    stk.push_back({d, dp[node]});
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent && !visited[neighbor]) {
            CalcDp(neighbor, node, d + w, stk);
        }
    }
    stk.pop_back();
}

void Add(int node, int parent) {
    se.insert({dp[node], tim++});
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent && !visited[neighbor]) {
            Add(neighbor, node);
        }
    }
}

void Erase(int node, int parent) {
    se.erase(se.lower_bound({dp[node], 0}));
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent && !visited[neighbor]) {
            Erase(neighbor, node);
        }
    }
}

ll dp2[MAX_N];
Seg seg;

void Solve(int node, int parent, int d1, int d2, vi &stk) {
    dp2[parent] = 0;
    if (d2 <= k) {
        ll cnt = se.order_of_key({k - d2 + 1, 0}) - se.order_of_key({k - d1 + 1, 0});
        dp2[parent] = cnt;
    }
    int r = lower_bound(all(stk), d1 - k + 1) - stk.begin();
    int l = lower_bound(all(stk), d2 - k) - stk.begin();
    dp2[parent] += seg.Query(l, r);

    ll s = real_sz[node] < real_sz[parent] ? real_sz[node] : (n - real_sz[parent]);
    ans[parent] += dp2[parent] * s;
    if (visited[node]) return;
    seg.Update(stk.size(), dp2[parent]);
    stk.push_back(d2);
    for (auto [neighbor, w] : tree[node]) {
        if (neighbor != parent) {
            Solve(neighbor, node, d1 + w, d1, stk);
        }
    }
    stk.pop_back();
}

void CentroidDecomposition(int node) {
    if (visited[node]) return;
    int centroid = FindCentroid(node);
    vii stk;
    CalcDp(centroid, -1, 0, stk);
//    for (auto [neighbor, w] : tree[centroid]) {
//        CalcDp(neighbor, node);
//    }
    dp[centroid] = 0;
    Add(centroid, -1);
    DfsSz(centroid, -1);
    for (auto [neighbor, w] : tree[centroid]) {
        if (!visited[neighbor]) Erase(neighbor, centroid);
        vi st;
        Solve(neighbor, centroid, w, 0, st);
        if (!visited[neighbor]) Add(neighbor, centroid);
    }
    se.clear();

    visited[centroid] = true;
    for (auto [neighbor, w] : tree[centroid]) {
        CentroidDecomposition(neighbor);
    }
}

int32_t main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> k;
    for (int i = 0; i < n - 1; i++) {
        int u, v, l;
        cin >> u >> v >> l;
        tree[u].push_back({v, l});
        tree[v].push_back({u, l});
    }
    seg = Seg(n + 1);
    RealDfsSz(0, -1);
    CentroidDecomposition(0);
    for (int i = 0; i < n; i++) cout << ans[i] << '\n';
}

Compilation message (stderr)

Main.cpp: In function 'int32_t main()':
Main.cpp:26:24: warning: '<anonymous>.Seg::n' is used uninitialized in this function [-Wuninitialized]
   26 |         seg.resize(2 * n);
      |                        ^

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