Submission #1144134

#	Time	Username	Problem	Language	Result	Execution time	Memory
1144134		Zero_OP	Petrol stations (CEOI24_stations)	C++17	100 / 100	385 ms	14704 KiB

Main

#include <bits/stdc++.h>

using namespace std;

#define mp make_pair
#define rep(i, l, r) for(int i = (l); i < (r); ++i)
#define all(v) begin(v), end(v)
#define rall(v) rbegin(v), rend(v)
#define sum_of(v) accumulate(all(v), 0ll)
#define compact(v) v.erase(unique(all(v)), end(v))
#define sz_of(v) (int)v.size()
#define dbg(x) "[" #x " = " << (x) << "]"
#define file(task) if(fopen(task".inp", "r")){ freopen(task".inp", "r", stdin); freopen(task".out", "w", stdout); }

template<typename T>
bool minimize(T& a, const T& b){
    if(a > b) return a = b, true;
    return false;
}

template<typename T>
bool maximize(T& a, const T& b){
    if(a < b) return a = b, true;
    return false;
}

using ll = long long;
using ull = unsigned long long;
using ld = long double;
using db = double;
using pi = pair<int, int>;
using pl = pair<ll, ll>;

using vi = vector<int>;
using vb = vector<bool>;
using vl = vector<ll>;
using vpi = vector<pi>;
using vpl = vector<pl>;

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

const int MAX = 7e4 + 5;

int N, K, sz[MAX], dp[MAX];
bool deleted[MAX];
ll sum[MAX], toward[MAX], result[MAX];
vpi adj[MAX];
vector<pair<ll, int>> cur;
vpl stck;

void reset_info(int u){
    dp[u] = 0;
    sum[u] = 0;
    toward[u] = 0;
}

int dfs_sz(int u, int p){
    reset_info(u);
    sz[u] = 1;
    for(auto [v, w] : adj[u]) if(v != p && !deleted[v]){
        sz[u] += dfs_sz(v, u);
    } return sz[u];
}

int find_centroid(int u, int p, int target){
    for(auto [v, w] : adj[u]) if(v != p && !deleted[v] && sz[v] * 2 > target){
        return find_centroid(v, u, target);
    } return u;
}

int find_stop(ll s){
    int l = 0, r = sz_of(cur) - 1, ans = -1;
    while(l <= r){
        int mid = l + r >> 1;
        if(s - cur[mid].first > K) ans = cur[mid + 1].second, l = mid + 1;
        else r = mid - 1;
    } return ans;
}

void dfs1(int u, int p, int extra, vpl& vec){
    cur.emplace_back(sum[u], u);
    int x = find_stop(sum[u]);
    for(auto [v, w] : adj[u]) if(v != p && !deleted[v]){
        sum[v] = sum[u] + w;
        dfs1(v, u, extra, vec);
    }
    cur.pop_back();

    result[u] += 1LL * extra * dp[u];
    if(x == -1) vec.emplace_back(sum[u], dp[u] + 1);
    else dp[x] += dp[u] + 1;
}

void dfs2(int u, int p, vpl& sub, int sign){
    int L = lower_bound(all(sub), mp(K - sum[u] + 1, 0ll)) - sub.begin();
    int R = upper_bound(all(sub), mp(K - sum[p] + 1, 0ll)) - sub.begin() - 1;
    ll cur = 0;
    if(L <= R) cur = sub[R].second - (L ? sub[L - 1].second : 0);

    // cout << dbg(u) << dbg(cur) << dbg(sign) << '\n';

    toward[u] += cur * sign;
    result[p] += 1LL * sz[u] * cur * sign;

    // if(p == 1) {
    //     for(auto [x, y] : sub) cout << x << ' ' << y << '\n';
    //     cout << dbg(L) << dbg(R) << '\n';
    //     cout << dbg(K - sum[u]) << dbg(K - sum[p]) << '\n';
    //     cout << dbg(result[1]) << '\n';
    // }

    for(auto [v, w] : adj[u]) if(v != p && !deleted[v]){
        dfs2(v, u, sub, sign);
    }
}

void dfs3(int u, int p){
    int L = lower_bound(all(stck), mp(sum[p] - K, 0ll)) - stck.begin();
    int R = upper_bound(all(stck), mp(sum[u] - K, 0ll)) - stck.begin() - 1;
    
    ll val = 0;
    if(L <= R){
        val = stck[R].second - (L ? stck[L - 1].second : 0);
        result[p] += 1LL * sz[u] * val;
        toward[u] += val;
    }

    stck.push_back(mp(sum[p], (stck.empty() ? 0 : stck.back().second) + toward[u]));

    for(auto [v, w] : adj[u]) if(v != p && !deleted[v]){
        dfs3(v, u);
    }

    stck.pop_back();
}

void solve(int u){
    u = find_centroid(u, -1, dfs_sz(u, -1)); dfs_sz(u, -1);
    sum[u] = 0; deleted[u] = true;
    vpl all_sub = {{0, 1}};
    cur = {{0, u}};
    for(auto [v, w] : adj[u]) if(!deleted[v]){
        sum[v] = sum[u] + w;
        vpl sub;
        dfs1(v, u, sz[u] - sz[v], sub);

        sort(all(sub));
        all_sub.insert(all_sub.end(), all(sub));
        rep(i, 1, sz_of(sub)) sub[i].second += sub[i - 1].second;
        dfs2(v, u, sub, -1);
    }

    sort(all(all_sub));
    rep(i, 1, sz_of(all_sub)) all_sub[i].second += all_sub[i - 1].second;

    for(auto [v, w] : adj[u]) if(!deleted[v]) dfs2(v, u, all_sub, +1);
    for(auto [v, w] : adj[u]) if(!deleted[v]) dfs3(v, u);
    for(auto [v, w] : adj[u]) if(!deleted[v]) solve(v);
}

void testcase(){
    cin >> N >> K;
    rep(i, 1, N){
        int u, v, l;
        cin >> u >> v >> l;
        adj[u].emplace_back(v, l);
        adj[v].emplace_back(u, l);
    }

    solve(0);
    rep(i, 0, N) cout << result[i] << '\n';
}

int main(){
    ios_base::sync_with_stdio(0); cin.tie(0);

#ifdef LOCAL
    freopen("in.txt", "r", stdin);
#endif // LOCAL

    int T = 1; //cin >> T;
    while(T--){
        testcase();
    }

    return 0;
}

• 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