Submission #1138806

#	Time	Username	Problem	Language	Result	Execution time	Memory
1138806		Zbyszek99	Road Closures (APIO21_roads)	C++20	100 / 100	166 ms	50528 KiB

roads

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
#define ll long long
#define ld long double
#define ull unsigned long long
#define ff first
#define ss second
#define pii pair<int,int>
#define pll pair<long long, long long>
#define vi vector<int>
#define vl vector<long long>
#define pb push_back
#define rep(i, b) for(int i = 0; i < (b); ++i)
#define rep2(i,a,b) for(int i = a; i <= (b); ++i)
#define rep3(i,a,b,c) for(int i = a; i <= (b); i+=c)
#define count_bits(x) __builtin_popcountll((x))
#define all(x) (x).begin(),(x).end()
#define siz(x) (int)(x).size()
#define forall(it,x) for(auto& it:(x))
using namespace __gnu_pbds;
using namespace std;
typedef tree<int, null_type, less<int>, rb_tree_tag,tree_order_statistics_node_update> ordered_set;
//mt19937 mt;void random_start(){mt.seed(chrono::time_point_cast<chrono::milliseconds>(chrono::high_resolution_clock::now()).time_since_epoch().count());}
//ll los(ll a, ll b) {return a + (mt() % (b-a+1));}
const int INF = 1e9+50;
const ll INF_L = 1e18+40;
const ll MOD = 1e9+7;

struct SEG_TREE
{
    int tree_siz = 0;
    pll* drzewo;
    ll get_kth_sum(int akt, int p1, int p2, int k)
    {
        if(k <= 0) return 0;
        if(tree_siz == 0) return -1;
       // cout << tree_siz << " " << akt << " " << p1 << " " << p2 << " " << drzewo[akt-1].ff << " " << drzewo[akt-1].ss << " " << k << " get\n";
        
        if(drzewo[akt-1].ss == k) return drzewo[akt-1].ff;
        if(drzewo[akt-1].ss < k) return -1;
        if(drzewo[akt*2-1].ss >= k) return get_kth_sum(akt*2,p1,(p1+p2)/2,k);
        return get_kth_sum(akt*2+1,(p1+p2)/2+1,p2,k - drzewo[akt*2-1].ss) + drzewo[akt*2-1].ff;
    }
    void upd(int v)
    {
        drzewo[v-1].ff = drzewo[v*2-1].ff + drzewo[v*2].ff;
        drzewo[v-1].ss = drzewo[v*2-1].ss + drzewo[v*2].ss;
        if(v != 1) upd(v/2);
    }
    void change(int ind, ll x, int is)
    {
        drzewo[tree_siz/2+ind] = {x,is};
        upd((tree_siz/2+ind+1)/2);
    }
    void set_siz(int N)
    {
        if(N == 0) return;
        tree_siz = (1 << (__lg(N)+2))-1;
        drzewo = new pll[tree_siz];
        for(int i = 0; i < tree_siz; i++)
        {
            drzewo[i] = {0,0};
        }
    }
};

SEG_TREE vtree[100001];
vector<pll> start_graph[100001];
set<int> graph[100001];
int parent[100001];
int deg[100001];
int index_in_parent[100001];
ll edge_val[100001];
int pre[100001];
ll dp[100001][2];
ll ans[100001];
bool used[100001];
int cur_pre = 0;

void dfs(int v, int pop)
{
    pre[v] = cur_pre++;
    vector<pii> pom;
    forall(it,start_graph[v])
    {
        if(it.ff != pop)
        {
            parent[it.ff] = v;
            edge_val[it.ff] = it.ss;
            dfs(it.ff,v);
            pom.pb({it.ss,it.ff});
        }
    }
    sort(all(pom));
    vtree[v].set_siz(siz(pom));
    rep(i,siz(pom))
    {
        index_in_parent[pom[i].ss] = i;
    }
}


vl minimum_closure_costs(int n, vi u, vi v, vi w)
{
    //random_start();
    edge_val[1] = 1e15;
    rep(i,n-1)
    {
        int a = u[i],b = v[i],c = w[i];
        a++;
        b++;
        start_graph[a].pb({b,c});
        start_graph[b].pb({a,c});
        deg[a]++;
        deg[b]++;
    }
    vector<pii> vert;
    dfs(1,1);
    rep2(i,1,n)
    {
        vert.pb({deg[i],i});
    }
    sort(all(vert));
    reverse(all(vert));
    rep2(i,2,n)
    {
        vtree[parent[i]].change(index_in_parent[i],edge_val[i],1);
    }
    set<pii> cur_vert;
    int ptr = 0;
    for(int k = n-1; k >= 0; k--)
    {
        while(ptr < siz(vert) && vert[ptr].ff > k)
        {
            int v= vert[ptr].ss;
            cur_vert.insert({-pre[v],v});
            forall(it,start_graph[v])
            {
                if(deg[it.ff] <= k || parent[v] == it.ff) continue;
                graph[v].insert(it.ff);
            }
            if(v != 1)
            {
                if(used[parent[v]])
                graph[parent[v]].insert(v);
                vtree[parent[v]].change(index_in_parent[v],0,0);
            }
            used[v] = 1;
            ptr++;
        }
        //cout << k << " k\n";
        forall(it2,cur_vert)
        {
            int v = it2.ss;
        //   if(k == 1) cout << v << " " << deg[v] << " v\n";
            dp[v][0] = 1e15;
            dp[v][1] = 1e15;
            ll cur_ans = 0;
            vector<ll> dif;
            forall(it,graph[v])
            {
                //cout << it << " graph\n";
                cur_ans += dp[it][1];
                dif.pb(dp[it][0]-dp[it][1]);
            }
            sort(all(dif));
            int zost = deg[v];
            rep(i,siz(dif)+1)
            {
            //   if(k == 0) cout << i << " " << cur_ans << " dif\n";
                ll kth1 = vtree[v].get_kth_sum(1,0,vtree[v].tree_siz/2,zost - (k+1));
            //   if(k == 0) cout << kth1 << " kth1\n";
                ll kth2 = vtree[v].get_kth_sum(1,0,vtree[v].tree_siz/2,zost - k);
         //     if(k == 0) cout << kth2 << " kth2\n";
                if(kth1 != -1) dp[v][0] = min(dp[v][0],cur_ans + kth1 + edge_val[v]);
                if(kth2 != -1) dp[v][1] = min(dp[v][1],cur_ans + kth2);
                if(i != siz(dif))
                {
                    cur_ans += dif[i];
                    zost--;
                }
            }
          //  if(k == 0) cout << dp[v][0] << " " << dp[v][1] << " " << edge_val[v] << " dp\n";
            if(deg[parent[v]] <= k) ans[k] += min(dp[v][0],dp[v][1]);
        }
    }
    vl ans2;
    rep2(k,0,n-1) ans2.pb(ans[k]);
    return ans2;
}

• 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