Submission #743472

#	Time	Username	Problem	Language	Result	Execution time	Memory
743472		Jarif_Rahman	Road Closures (APIO21_roads)	C++17	100 / 100	438 ms	37428 KiB

roads

#include "roads.h"
#include <bits/stdc++.h>
#define pb push_back
#define f first
#define sc second
using namespace std;
typedef long long int ll;
typedef string str;

namespace Treap{
    mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
    struct Node{
        ll X, Y;
        Node *l = nullptr, *r = nullptr;

        int sz = 1;
        ll sum = 0;

        Node(){}
        Node(ll _X): X(_X), Y(uniform_int_distribution<ll>(0, INT_MAX-1)(rng)), sum(_X){}
        void update_node(){
            sz = 1, sum = X;
            if(l) sz+=l->sz, sum+=l->sum;
            if(r) sz+=r->sz, sum+=r->sum;
        }
    };
    int sz(Node* nd){
        if(!nd) return 0;
        return nd->sz;
    }

    pair<Node*, Node*> split_by_x(Node* root, ll X){
        if(!root) return {nullptr, nullptr};
        if(root->X <= X){
            auto [L, R] = split_by_x(root->r, X);
            root->r = L;
            root->update_node();
            return {root, R};
        }
        else{
            auto [L, R] = split_by_x(root->l, X);
            root->l = R;
            root->update_node();
            return {L, root};
        }
    }
    pair<Node*, Node*> split_by_sz(Node* root, int i){
        if(!root) return {nullptr, nullptr};
        if(sz(root->l)+1 <= i){
            auto [L, R] = split_by_sz(root->r, i-sz(root->l)-1);
            root->r = L;
            root->update_node();
            return {root, R};
        }
        else{
            auto [L, R] = split_by_sz(root->l, i);
            root->l = R;
            root->update_node();
            return {L, root};
        }
    }

    Node* merge(Node* L, Node* R){
        if(!L) return R;
        if(!R) return L;
        if(L->Y > R->Y){
            L->r = merge(L->r, R);
            L->update_node();
            return L;
        }
        else{
            R->l = merge(L, R->l);
            R->update_node();
            return R;
        }
    }

    void insert(Node* &root, ll X){
        auto [A, B] = split_by_x(root, X);
        A = merge(A, new Node(X));
        root = merge(A, B);
    }
    void erase(Node* &root, ll X){
        auto [A, B] = split_by_x(root, X-1);
        auto [C, D] = split_by_sz(B, 1);
        if(C) delete C;
        root = merge(A, D);
    }
    ll sum(Node* &root, int i){
        auto [A, B] = split_by_sz(root, i);
        ll rt = A->sum;
        root = merge(A, B);
        return rt;
    }
};

int n, k;
vector<vector<pair<int, int>>> tree;
vector<int> p;
vector<ll> pe;
vector<int> deg;
vector<int> depth;

vector<ll> dp[2];
vector<vector<int>> forest;
vector<bool> visited;
vector<Treap::Node*> sth;

void pre_dfs(int nd, int ss, int d = 0){
    depth[nd] = d;
    for(auto [x, w]: tree[nd]) if(x != ss){
        pre_dfs(x, nd, d+1);
        p[x] = nd;
        pe[x] = w;
        Treap::insert(sth[nd], w);
    }
}

void dfs(int nd){
    visited[nd] = 1;
    ll sum = 0;
    for(int x: forest[nd]){
        dfs(x);
        sum+=dp[1][x];
        Treap::insert(sth[nd], -dp[1][x]+dp[0][x]+pe[x]);
    }

    dp[0][nd] = sum, dp[1][nd] = sum;

    int c = sth[nd]->sz;
    if(c-k > 0){
        dp[0][nd]+=Treap::sum(sth[nd], c-k);
    }
    if(c-k+1 > 0){
        dp[1][nd]+=Treap::sum(sth[nd], c-k+1);
    }

    for(int x: forest[nd]){
        Treap::erase(sth[nd], -dp[1][x]+dp[0][x]+pe[x]);
    }

    dp[1][nd] = min(dp[1][nd], dp[0][nd]+pe[nd]);
}

vector<ll> minimum_closure_costs(int _n, vector<int> U, vector<int> V, vector<int> W){
    swap(n, _n);
    tree.resize(n);
    p.assign(n, -1);
    pe.assign(n, 0);
    deg.assign(n, 0);
    depth.resize(n);

    fill(dp, dp+2, vector<ll>(n, 0));
    forest.resize(n);
    visited.assign(n, 0);
    sth.assign(n, nullptr);

    for(int i = 0; i < n-1; i++){
        deg[U[i]]++, deg[V[i]]++;
        tree[U[i]].pb({V[i], W[i]});
        tree[V[i]].pb({U[i], W[i]});
    }

    pre_dfs(0, -1);

    vector<int> o(n);
    for(int i = 0; i < n; i++) o[i] = i;
    sort(o.begin(), o.end(), [&](int a, int b){
        return deg[a] > deg[b];
    });

    vector<ll> ans(n, 0);
    for(int w: W) ans[0]+=w;
    for(k = n-1; k > 0; k--){
        vector<int> nodes;
        for(int i = 0; i < n; i++){
            if(deg[o[i]] <= k) break;
            nodes.pb(o[i]);
            if(deg[p[o[i]]] > k){
                forest[p[o[i]]].pb(o[i]);
                Treap::erase(sth[p[o[i]]], pe[o[i]]);
            }
        }

        sort(nodes.begin(), nodes.end(), [&](int a, int b){
            return depth[a] < depth[b];
        });

        for(int nd: nodes) if(!visited[nd]) dfs(nd), ans[k]+=dp[1][nd];

        for(int nd: nodes){
            forest[nd].clear();
            visited[nd] = 0;
            if(deg[p[nd]] > k){
                Treap::insert(sth[p[nd]], pe[nd]);
            }
        }
    }

    return ans;
}

• 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