Submission #299867

# Submission time Handle Problem Language Result Execution time Memory
299867 2020-09-15T22:16:25 Z caoash Transport (COCI19_transport) C++14
78 / 130
1000 ms 37744 KB
#include <bits/stdc++.h> 
using namespace std;

using ll = long long;

using vi = vector<int>;
using vl = vector<ll>;
#define pb push_back
#define rsz resize
#define all(x) begin(x), end(x)
#define sz(x) (int)(x).size()

using pi = pair<int,int>;
#define f first
#define s second
#define mp make_pair

const int MX = 200005;
const int MOD = (int) (1e9 + 7);
const ll INF = (ll) 1e18;

namespace output {
    void pr(int x) { cout << x; }
    void pr(long x) { cout << x; }
    void pr(ll x) { cout << x; }
    void pr(unsigned x) { cout << x; }
    void pr(unsigned long x) { cout << x; }
    void pr(unsigned long long x) { cout << x; }
    void pr(float x) { cout << x; }
    void pr(double x) { cout << x; }
    void pr(long double x) { cout << x; }
    void pr(char x) { cout << x; }
    void pr(const char* x) { cout << x; }
    void pr(const string& x) { cout << x; }
    void pr(bool x) { pr(x ? "true" : "false"); }
    
    template<class T1, class T2> void pr(const pair<T1,T2>& x);
    template<class T> void pr(const T& x);
    
    template<class T, class... Ts> void pr(const T& t, const Ts&... ts) { 
        pr(t); pr(ts...); 
    }
    template<class T1, class T2> void pr(const pair<T1,T2>& x) { 
        pr("{",x.f,", ",x.s,"}"); 
    }
    template<class T> void pr(const T& x) { 
        pr("{"); // const iterator needed for vector<bool>
        bool fst = 1; for (const auto& a: x) pr(!fst?", ":"",a), fst = 0; 
        pr("}");
    }
    
    void ps() { pr("\n"); } // print w/ spaces
    template<class T, class... Ts> void ps(const T& t, const Ts&... ts) { 
        pr(t); if (sizeof...(ts)) pr(" "); ps(ts...); 
    }
    
    void pc() { cout << "]" << endl; } // debug w/ commas
    template<class T, class... Ts> void pc(const T& t, const Ts&... ts) { 
        pr(t); if (sizeof...(ts)) pr(", "); pc(ts...); 
    }
    #define dbg(x...) pr("[",#x,"] = ["), pc(x);
}
 
#ifdef LOCAL
using namespace output;
#endif

#include <ext/pb_ds/tree_policy.hpp>
#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;
template <class T> using Tree = tree<T, null_type, less<T>, 
    rb_tree_tag, tree_order_statistics_node_update>; 

int n;
ll a[MX];
bool blocked[MX];
vector<pair<ll, int>> adj[MX];
ll ans = 0;

Tree<pair<ll, int>> vals[MX];
ll dist[MX];
ll mdist[MX];
int sz[MX];

void dfs(int v, int p) {
    if (p != -1 && mdist[v] >= 0) ++ans;
    dist[v] += a[v];
    for (pi to : adj[v]) {
        if (blocked[to.f]) continue;
        if (to.f != p) {
            dist[to.f] = dist[v] - to.s;
            mdist[to.f] = min(mdist[v], dist[to.f]);
            dfs(to.f, v);
        }
    }
}

ll need = 0;

void comb(Tree<pair<ll, int>> &a, Tree<pair<ll, int>> &b) {
    if (sz(a) < sz(b)) swap(a, b);
    for (auto x : b) {
        a.insert(x);
        ++need;
    }
    b.clear();
}

void dfs2(int v, int p) {
    if (p != -1) vals[v].insert(mp(dist[v], v));
    for (pi to : adj[v]) {
        if (blocked[to.f]) continue;
        if (to.f != p) {
            dfs2(to.f, v);
            comb(vals[v], vals[to.f]);
        }
    }
    // dbg(v, vals[v]);
    // dbg(sz(vals[v]));
    int rem = vals[v].order_of_key(mp(dist[v], INT_MIN));
    for (int i = 0; i < rem; i++) {
        ++need;
        vals[v].erase(vals[v].begin());
    }
    // dbg(v, vals[v]);
}

vi order;

void dfs3(int v, int p) {
    for (pi to : adj[v]) {
        if (blocked[to.f]) continue;
        if (to.f != p) {
            dfs3(to.f, v);
        }
    }
    order.pb(v);
}

void solve(int root) {
    dfs3(root, -1);
    for (int x : order) mdist[x] = INF;
    // dbg(root, sz(order));
    mdist[root] = 0;
    dist[root] = 0;
    dfs(root, -1);
    dfs2(root, -1);
    Tree<pair<ll, int>> &fin = vals[root];
    vector<bool> ok(n);
    for (auto &x : fin) {
        ok[x.s] = true;
    }
    // dbg(fin);
    ans += sz(fin);
    order.clear();
    for (pi to : adj[root]) {
        if (blocked[to.f]) continue;
        dfs3(to.f, root);
        for (int x : order) {
            if (ok[x]) {
                fin.erase(mp(dist[x], x));     
            }
        }
        // dbg(to.f, fin);
        for (int x : order) {
            ll cdist = mdist[x];
            // dbg(x, cdist);
            // dbg(x, -cdist + a[root]);
            // dbg(ans);
            ans += sz(fin) - fin.order_of_key(mp(-cdist + a[root], INT_MIN));
            // dbg(ans);
        }
        for (int x : order) {
            if(ok[x]) fin.insert(mp(dist[x], x)); 
        }
        order.clear();
    }
    fin.clear();
    // ps(ans);
}

void dfs4(int v, int p) {
    sz[v] = 1;
    for (pi to : adj[v]) {
        if (blocked[to.f]) continue;
        if (to.f != p) {
            dfs4(to.f, v);
            sz[v] += sz[to.f];
        }
    }
}

int find(int v, int p, int root) {
    // dbg("FINDING", v);
    for (pi to : adj[v]) {
        if (blocked[to.f]) continue;
        if (to.f != p) {
            if (sz[to.f] > sz[root] / 2) {
                return find(to.f, v, root);
            }
        }
    }
    return v;
}

void go(int root) {
    dfs4(root, -1);
    int cc = find(root, -1, root);
    // dbg(root, cc);
    solve(cc);
    // dbg(cc, need);
    // dbg(cc, ans);
    blocked[cc] = true;
    for (pi to : adj[cc]) {
        if (!blocked[to.f]) {
            go(to.f);
        }
    }
}

int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin >> n;
    for (int i = 0; i < n; i++) cin >> a[i];
    for (int i = 0; i < n - 1; i++) {
        int u, v, w; cin >> u >> v >> w;
        u--, v--;
        adj[u].pb(mp(v, w)), adj[v].pb(mp(u, w));
    }
    // solve(1);
    // dbg(need);
    //dbg("KEKW");
    go(0);
    cout << ans << '\n';
} 
# Verdict Execution time Memory Grader output
1 Correct 42 ms 24192 KB Output is correct
2 Correct 37 ms 24320 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 236 ms 24696 KB Output is correct
2 Correct 60 ms 24832 KB Output is correct
# Verdict Execution time Memory Grader output
1 Execution timed out 1091 ms 31224 KB Time limit exceeded
2 Halted 0 ms 0 KB -
# Verdict Execution time Memory Grader output
1 Execution timed out 1083 ms 34168 KB Time limit exceeded
2 Halted 0 ms 0 KB -
# Verdict Execution time Memory Grader output
1 Execution timed out 1092 ms 37744 KB Time limit exceeded
2 Halted 0 ms 0 KB -
# Verdict Execution time Memory Grader output
1 Correct 696 ms 28016 KB Output is correct
2 Correct 150 ms 27000 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 300 ms 29304 KB Output is correct
2 Correct 475 ms 27768 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 729 ms 30200 KB Output is correct
2 Correct 519 ms 31348 KB Output is correct
# Verdict Execution time Memory Grader output
1 Correct 970 ms 32240 KB Output is correct
2 Correct 835 ms 33524 KB Output is correct
# Verdict Execution time Memory Grader output
1 Execution timed out 1102 ms 33776 KB Time limit exceeded
2 Halted 0 ms 0 KB -