Submission #1114872

#	Time	Username	Problem	Language	Result	Execution time	Memory
1114872		vladilius	Road Closures (APIO21_roads)	C++17	100 / 100	884 ms	368616 KiB

roads

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
using pil = pair<int, ll>;
using pll = pair<ll, ll>;
#define pb push_back
#define ff first
#define ss second
#define arr3 array<ll, 3>

struct DS{
    struct node{
        node *l, *r;
        int c;
        ll s;
        node(){
            l = r = 0;
            c = s = 0;
        }
    };
    ll n;
    node *rt;
    void init(){
        n = 1e14;
        rt = new node();
    }
    void add(node *v, ll tl, ll tr, ll& p){
        if (tl == tr){
            v -> c++;
            v -> s += p;
            return;
        }
        ll tm = (tl + tr) / 2;
        if (p <= tm){
            if (!(v -> l)) v -> l = new node();
            add(v -> l, tl, tm, p);
        }
        else {
            if (!(v -> r)) v -> r = new node();
            add(v -> r, tm + 1, tr, p);
        }
        v -> c = 0;
        if (v -> l) v -> c += v -> l -> c;
        if (v -> r) v -> c += v -> r -> c;
        v -> s = 0;
        if (v -> l) v -> s += v -> l -> s;
        if (v -> r) v -> s += v -> r -> s;
    }
    void add(ll x){
        add(rt, 1, n, x);
    }
    void rem(node *v, ll tl, ll tr, ll& p){
        if (tl == tr){
            v -> c--;
            v -> s -= p;
            return;
        }
        ll tm = (tl + tr) / 2;
        if (p <= tm){
            rem(v -> l, tl, tm, p);
        }
        else {
            rem(v -> r, tm + 1, tr, p);
        }
        v -> c = 0;
        if (v -> l) v -> c += v -> l -> c;
        if (v -> r) v -> c += v -> r -> c;
        v -> s = 0;
        if (v -> l) v -> s += v -> l -> s;
        if (v -> r) v -> s += v -> r -> s;
    }
    void rem(ll x){
        rem(rt, 1, n, x);
    }
    int sz(){
        return rt -> c;
    }
    ll fn(node *v, ll tl, ll tr, int k){
        if (tl == tr) return tl;
        ll tm = (tl + tr) / 2;
        if (v -> r && (v -> r -> c >= k)){
            return fn(v -> r, tm + 1, tr, k);
        }
        int S = (v -> r) ? v -> r -> c : 0;
        return fn(v -> l, tl, tm, k - S);
    }
    pil sum(node *v, ll tl, ll tr, ll l, ll r){
        if (l > tr || r < tl) return {0, 0};
        if (l <= tl && tr <= r) return {v -> c, v -> s};
        ll tm = (tl + tr) / 2;
        pil out;
        if (v -> l){
            auto [c, s] = sum(v -> l, tl, tm, l, r);
            out.ff += c; out.ss += s;
        }
        if (v -> r){
            auto [c, s] = sum(v -> r, tm + 1, tr, l, r);
            out.ff += c; out.ss += s;
        }
        return out;
    }
    ll get(int k){
        k = min(k, sz());
        if (!k) return 0;
        ll p = fn(rt, 1, n, k);
        if (p == n) return n * k;
        auto [c, s] = sum(rt, 1, n, p + 1, n);
        return s + p * (k - c);
    }
    void clear(){
        rt = new node();
    }
};
 
vector<ll> minimum_closure_costs(int n, vector<int> U, vector<int> V, vector<int> W){
    vector<pii> g[n + 1];
    ll tot = 0;
    for (int i = 0; i < n - 1; i++){
        int u = U[i] + 1, v = V[i] + 1, w = W[i];
        g[u].pb({v, w});
        g[v].pb({u, w});
        tot += w;
    }
    
    vector<pil> ch1[n + 1], ch2[n + 1];
    vector<pll> p(n + 1);
    vector<ll> W1(n + 1);
    DS T; T.init();
    function<void(int, int)> dfs = [&](int v, int pr){
        if ((g[v].size() + !pr) == 1){
            ch1[v] = {{0, 0}};
            ch2[v] = {{0, 0}};
            return;
        }
        for (auto [i, w]: g[v]){
            if (i == pr) continue;
            dfs(i, v);
        }
        
        vector<arr3> t;
        T.clear();
        for (auto [i, w]: g[v]){
            if (i == pr) continue;
            p[i].ff = p[i].ss = 0;
            T.add(w);
            W1[i] = w;
            for (auto [j, x]: ch1[i]){
                t.pb({j, i, x});
            }
            for (auto [j, x]: ch2[i]){
                t.pb({j, -i, x});
            }
        }
        
        sort(t.begin(), t.end());
        int p1 = 0; ll sum = 0;
        while (p1 < t.size()){
            int p2 = p1;
            while (p2 < t.size() && t[p1][0] == t[p2][0]){
                auto [j, i, x] = t[p2];
                if (i > 0){
                    ll f = p[i].ff + W1[i] - p[i].ss;
                    if (f > 0) T.rem(f);
                    p[i].ff = x;
                    f = p[i].ff + W1[i] - p[i].ss;
                    if (f > 0) T.add(f);
                }
                else {
                    i = -i;
                    sum -= p[i].ss;
                    ll f = p[i].ff + W1[i] - p[i].ss;
                    if (f > 0) T.rem(f);
                    p[i].ss = x;
                    f = p[i].ff + W1[i] - p[i].ss;
                    if (f > 0) T.add(f);
                    sum += p[i].ss;
                }
                p2++;
            }
            
            ch2[v].pb({t[p1][0], sum + T.get((int) t[p1][0])});
            if (!t[p1][0]) ch1[v].pb({0, sum});
            else ch1[v].pb({t[p1][0], sum + T.get((int) t[p1][0] - 1)});
 
            int R = (int)((p2 == (int) t.size()) ? n : t[p2][0]);
            for (int j = (int) t[p1][0] + 1; j <= min(R - 1, T.sz()); j++){
                ch2[v].pb({j, sum + T.get(j)});
            }
            for (int j = (int) t[p1][0] + 1; j <= min(R - 1, T.sz() + 1); j++){
                if (!j) ch1[v].pb({0, sum});
                else ch1[v].pb({j, sum + T.get(j - 1)});
            }
            p1 = p2;
        }
    };
    dfs(1, 0);
    
    vector<ll> dp(n);
    for (int i = 0; i < ch2[1].size(); i++){
        int R = (i == (ch2[1].size() - 1)) ? n : ch2[1][i + 1].ff;
        for (int j = ch2[1][i].ff; j < R; j++){
            dp[j] = ch2[1][i].ss;
        }
    }
    
    vector<ll> out;
    for (int i = 0; i < n; i++){
        out.pb(tot - dp[i]);
    }
    return out;
}

Compilation message (stderr)

roads.cpp: In lambda function:
roads.cpp:158:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  158 |         while (p1 < t.size()){
      |                ~~~^~~~~~~~~~
roads.cpp:160:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::array<long long int, 3> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  160 |             while (p2 < t.size() && t[p1][0] == t[p2][0]){
      |                    ~~~^~~~~~~~~~
roads.cpp: In function 'std::vector<long long int> minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)':
roads.cpp:200:23: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  200 |     for (int i = 0; i < ch2[1].size(); i++){
      |                     ~~^~~~~~~~~~~~~~~
roads.cpp:201:20: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  201 |         int R = (i == (ch2[1].size() - 1)) ? n : ch2[1][i + 1].ff;
      |                  ~~^~~~~~~~~~~~~~~~~~~~~~

• Subtask 1: 5 / 5
• Subtask 2: 7 / 7
• Subtask 3: 14 / 14
• Subtask 4: 10 / 10
• Subtask 5: 17 / 17
• Subtask 6: 25 / 25
• Subtask 7: 22 / 22