Submission #1075856

#	Time	Username	Problem	Language	Result	Execution time	Memory
1075856		ProtonDecay314	Petrol stations (CEOI24_stations)	C++17	18 / 100	103 ms	33724 KiB

Main

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef vector<ll> vll;
typedef vector<vll> vvll;
typedef vector<int> vi;
typedef vector<vi> vvi;
typedef pair<int, int> pi;
typedef vector<pi> vpi;
typedef vector<vpi> vvpi;
typedef pair<ll, ll> pll;
typedef vector<pll> vpll;
typedef vector<vpll> vvpll;
typedef vector<bool> vb;
typedef set<ll> sll;
#define IOS cin.tie(0); cout.tie(0); ios_base::sync_with_stdio(false)
#define INF(dtype) numeric_limits<dtype>::max()
#define NINF(dtype) numeric_limits<dtype>::min()
#define fi first
#define se second
#define pb push_back

typedef vector<vb> vvb;
typedef vector<vvb> v3b;
typedef vector<v3b> v4b;
typedef vector<string> vs;

// struct state {
//     int i, p;
//     int fuel;
// };

// void compute_depth(int i, int d, const vvi& adj, vi& depth) {
//     depth[i] = d;

//     // cout << i << " " << d << endl;

//     for(int j : adj[i]) {
//         compute_depth(j, d + 1, adj, depth);
//     }
// }

// int compute_size(int i, const vvi& adj, vi& tree_sz) {
//     int& ans = tree_sz[i];

//     ans += 1;

//     for(int j : adj[i]) {
//         ans += compute_size(j, adj, tree_sz);
//     }

//     return ans;
// }

ll compute_size(ll i, ll p, const vvll& adj, vll& tree_sz) {
    ll& ans = tree_sz[i];

    if(ans == -1ll) {
        ans = 1ll;

        for(ll j : adj[i]) {
            if(j == p) continue;
            ans += compute_size(j, i, adj, tree_sz);
        }
    }

    return ans;
}

ll solve(ll i, ll n, ll k) {
    ll nl = i / k;
    ll nr = (n - i - 1) / k;

    return nl * (n - i - 1) + nr * i;
}

vll solve(ll n, ll k, const vvpll& adj) {
    vll orig_ind(n, 0ll); // new index -> old index
    vll new_ind(n, 0ll); // old index -> new index
    vll lens(n - 1ll, 0ll);
    vll len_pref(n, 0ll);
    
    // Find the starting point (i.e., the point of outdegree 1)
    ll s = 0;

    for(ll i = 0; i < n; i++) {
        if(adj[i].size() == 1) {
            s = i;
            break;
        }
    }

    // ll ind = 0;
    orig_ind[0ll] = s;
    new_ind[s] = 0ll;

    ll prev = s;
    for(ll i = 0; i < n - 1; i++) {
        // Goto next
        // cout << s << " " << prev << endl;
        if(adj[s][0].fi != prev) {
            lens[i] = adj[s][0].se;
            prev = s;
            s = adj[s][0].fi;
        } else {
            lens[i] = adj[s][1].se;
            prev = s;
            s = adj[s][1].fi;
        }

        len_pref[i + 1ll] = len_pref[i] + lens[i];
        // cout << len_pref[i + 1ll] << " " << lens[i] << endl;

        // ind++;

        orig_ind[i + 1ll] = s;
        new_ind[s] = i + 1ll;

        // cout << s << " ";
    }
    // cout << endl;

    // Create new adjacency list

    vvpll adjn;
    for(ll i = 0; i < n; i++) {
        vpll adjnr;
        adjn.pb(adjnr);
    }

    for(ll i = 0ll; i < n; i++) {
        for(auto [j, l] : adj[i]) {
            adjn[new_ind[i]].pb({new_ind[j], l});
            adjn[new_ind[j]].pb({new_ind[i], l});
        }
    }

    // Create forward and backward graphs
    vvll adjf;
    vll f_indeg(n, 0ll);
    vvll adjb;
    vll b_indeg(n, 0ll);

    for(ll i = 0ll; i < n; i++) {
        vll adjfr;
        vll adjbr;

        adjf.pb(adjfr);
        adjb.pb(adjbr);
    }

    // Create forward graph (for each i, find the furthest "reach" to the right)
    for(ll i = 0ll; i < n; i++) {
        ll l = i, r = n;

        while(r - l > 1ll) {
            ll m = (l + r) >> 1ll;

            if(len_pref[m] - len_pref[i] <= k) l = m;
            else r = m;
        }

        // We stop at l
        // Edge case: l is n - 1, then we don't do anything

        if(l < n - 1ll) {
            adjf[l].pb(i);
            f_indeg[i]++;
        }
    }

    // Create backward graph (for each i, find the furthest "reach" to the left)
    for(ll i = 0ll; i < n; i++) {
        ll l = -1ll, r = i;

        while(r - l > 1ll) {
            ll m = (l + r) >> 1ll;

            if(len_pref[i] - len_pref[m] <= k) r = m;
            else l = m;
        }

        // We stop at l
        // Edge case: l is n - 1, then we don't do anything

        if(r > 0ll) {
            adjb[r].pb(i);
            b_indeg[i]++;
        }
    }

    // Compute subtree sizes for each i in the two graphs
    vll fsz(n, -1ll);
    vll bsz(n, -1ll);

    for(ll i = 0ll; i < n; i++) {
        if(f_indeg[i] == 0ll) {
            // This is a root in the forest, compute stuff
            compute_size(i, i, adjf, fsz);
        }
        if(b_indeg[i] == 0ll) {
            // This is a root in the forest, compute stuff
            compute_size(i, i, adjb, bsz);
        }
    }

    // Compute answers
    vll ans(n, 0ll);

    for(ll i = 0ll; i < n; i++) {
        ll ni = new_ind[i];

        ans[i] = (fsz[ni] - 1ll) * (n - ni - 1ll) + (bsz[ni] - 1ll) * (ni);
    }

    return ans;
}

int main() {
    ll n, k; cin >> n >> k;

    vvpll adj;
    for(ll i = 0; i < n; i++) {
        vpll adjr;
        adj.pb(adjr);
    }

    for(ll i = 0; i < n - 1; i++) {
        ll u, v, l;
        cin >> u >> v >> l;

        adj[u].pb({v, l});
        adj[v].pb({u, l});
    }

    vll ans = solve(n, k, adj);

    for(ll v : ans) cout << v << "\n";

    cout << flush;

    return 0;
}

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