Submission #361227

#	Time	Username	Problem	Language	Result	Execution time	Memory
361227		jainbot27	Transport (COCI19_transport)	C++17	0 / 130	984 ms	65540 KiB

Main

#include <bits/stdc++.h>
using namespace std;

#define f first
#define s second
#define pb push_back
#define ar array
#define all(x) x.begin(), x.end()
#define siz(x) (int)x.size()

#define FOR(x, y, z) for(int x = (y); x < (z); x++)
#define ROF(x, z, y) for(int x = (y-1); x >= (z); x--)
#define F0R(x, z) FOR(x, 0, z)
#define R0F(x, z) ROF(x, 0, z)
#define trav(x, y) for(auto&x:y)

using ll = long long;
using vi = vector<int>;
using vl = vector<long long>;
using pii = pair<int, int>;
using vpii = vector<pair<int, int>>;

template<class T> inline bool ckmin(T&a, T b) {return b < a ? a = b, 1 : 0;}
template<class T> inline bool ckmax(T&a, T b) {return b > a ? a = b, 1 : 0;}
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
//#define int long long
using namespace __gnu_pbds;
#define ordered_set tree<pair<ll, ll>, null_type,less<pair<ll, ll>>, rb_tree_tag,tree_order_statistics_node_update>

const char nl = '\n';
const int mxN = 2e5 + 10;
const int MOD = 1e9 + 7;
const long long infLL = 1e18;

struct centroid{
    vector<vector<int>> adj;
    vector<bool> vis;
    vi par, sz;
    int n;
    void init(int _n){ n = _n; adj.resize(n); par.resize(n); sz.resize(n); vis.resize(n, 0);}
    void add(int a, int b) {adj[a].pb(b); adj[b].pb(a);}
    int find_size(int u, int p = -1) {
        if(vis[u]) return 0;
        sz[u] = 1;
        trav(v, adj[u]){
            if(v != p)
                sz[u] += find_size(v, u);
        }
        return sz[u];
    }
    int find_centroid(int u, int p, int N){
        trav(v, adj[u]){
            if(v!=p&&!vis[v]&&sz[v]>N/2)
                return find_centroid(v, u, N);
        }
        return u;
    }
    void init_centroid(int u = 0, int p = -1){
        find_size(u);
        int c = find_centroid(u, -1, sz[u]);
        vis[c] = 1;
        par[c] = p;
        trav(v, adj[c])
            if(!vis[v])
                init_centroid(v, c);
    }
};

centroid C;
int n;
vector<pair<int, ll>> adj[mxN];
set<int> padj[mxN];
set<pair<int, int>> al;
set<int> nodes[mxN];
bool active[mxN];
ll cost[mxN];
ll a[mxN];
vl costs; vi to_remove;
ll ans = 0;

void dfs(int u, int p){
    //cout << u << ' ' << p << endl;
    trav(v, adj[u]){
        if(v.f == p || !active[v.f]) continue;
        cost[v.f] = cost[u] + a[u] - v.s;
        dfs(v.f, u);
    }
}

void dfs2(int u, int p, ll Cost, ll sum){
    to_remove.pb(u);
    //cout << "BRUH: " << u << ' ' << Cost << ' ' << sum << nl;
    if(Cost >= 0){
        costs.pb(sum);
    }
    trav(v, adj[u]){
        if(p == v.f || !active[v.f]) continue;
        if(active[v.f] && v.f != u){
            dfs2(v.f, u, min(a[v.f] - v.s, Cost + (a[v.f] - v.s)), sum + a[v.f] - v.s);
        }
    }
}

int32_t main(){
    ios_base::sync_with_stdio(0); cin.tie(0);
    cin >> n;
    F0R(i, n){
        cin >> a[i];
    }
    C.init(n);
    F0R(i, n-1){
        int e1, e2; ll w;
        cin >> e1 >> e2 >> w;
        e1--; e2--; 
        adj[e1].pb({e2, w});
        adj[e2].pb({e1, w});
        C.add(e1, e2);
    }
    C.init_centroid();
    F0R(i, n){
        if(C.par[i]!=-1)
            padj[C.par[i]].insert(i);
    }
    F0R(i, n){
        al.insert({siz(padj[i]), i});
        nodes[i].insert(i);
    }
    //cout << "WE GOT HERE" << endl;

    while(siz(al)){
        pii x = *al.begin();
        assert(x.f == 0) ;
        al.erase(al.begin());
        //cout << "FOR: " << x.s << endl;
        trav(X, nodes[x.s]){
            //cout << X << endl;
            active[X] = 1;
        }
        // do stuff here 

        cost[x.s] = 0;
        dfs(x.s, -1);

        ordered_set S;
        trav(X, nodes[x.s]){
            S.insert({cost[X], X});
            if(cost[X] >= 0) {
                //cout << "ADD: " << X << ' ' << x.s << nl;
                ans++;
            }
            //cout << X << ' ' << cost[X] << endl;
        }
        trav(c, adj[x.s]){
            int child = c.f; 
            if(!active[child]) continue;
            //cout << "WAIT A MINUTE BITCH" << endl;
            to_remove.clear(); costs.clear();
            dfs2(child, x.s, a[c.f]-c.s, a[c.f]-c.s);
            //ans += siz(costs);
            trav(X, to_remove){
                S.erase({cost[X], X});
            }
            trav(X, costs){
                ans += siz(S) - S.order_of_key({-X, -1});
                //cout << "CHECK " << siz(S) - S.order_of_key({-X, -1}) << ' ' << -X << endl;
            }
            //cout << "S:\n";
            //trav(x, S) cout << x.f << ' ';
            //cout << nl;
            trav(X, to_remove){
                S.insert({cost[X], X});
            }
        }

        trav(X, nodes[x.s]){
            active[X] = 0;
        }
        if(C.par[x.s] != -1){
            int p = C.par[x.s];
            al.erase({siz(padj[p]), p});
            padj[p].erase(x.s);
            al.insert({siz(padj[p]), p});
            if(siz(nodes[p]) < siz(nodes[x.s])){
                swap(nodes[p], nodes[x.s]);
            }
            trav(X, nodes[x.s]){
                nodes[p].insert(X);
            }
        }
    }
    cout << ans - n<< nl;
    return 0;
}

• Subtask 1: 0 / 13
• Subtask 2: 0 / 13
• Subtask 3: 0 / 13
• Subtask 4: 0 / 13
• Subtask 5: 0 / 13
• Subtask 6: 0 / 13
• Subtask 7: 0 / 13
• Subtask 8: 0 / 13
• Subtask 9: 0 / 13
• Subtask 10: 0 / 13