Submission #559294

#	Time	Username	Problem	Language	Result	Execution time	Memory
559294		nghiass001	Road Closures (APIO21_roads)	C++17	100 / 100	191 ms	36424 KiB

roads

#include <bits/stdc++.h>
#define FOR(i,l,r) for(int i=(l); i<=(r); ++i)
#define REP(i,l,r) for(int i=(l); i<(r); ++i)
#define FORD(i,r,l) for(int i=(r); i>=(l); --i)
#define REPD(i,r,l) for(int i=(r)-1; i>=(l); --i)
#define pii pair<int, int>
using namespace std;
mt19937 rd(chrono::steady_clock::now().time_since_epoch().count());
const int N = 1e5 + 5;

#define getsz(x) (x ? x->sz : 0)
#define getsum(x) (x ? x->sum : 0)
typedef struct node {
    int sz, cthis, priority;
    long long val, sum;
    node *L, *R;
    node(int _val) {
        L = R = nullptr;
        sz = cthis = sum = 0;
        val = _val;
        priority = rd() % 1000000000;
    }
    void rebuild() {
        sz = cthis + (L ? L->sz : 0) + (R ? R->sz : 0);
        sum = cthis * val + (L ? L->sum : 0) + (R ? R->sum : 0);
    }
} *TNode;

struct Treap {
    TNode root;

    TNode rotate_left(TNode x) {
        TNode y = x->R;
        TNode z = y->L;
        y->L = x;
        x->R = z;

        x->rebuild();
        y->rebuild();

        return y;
    }

    TNode rotate_right(TNode x) {
        TNode y = x->L;
        TNode z = y->R;
        y->R = x;
        x->L = z;

        x->rebuild();
        y->rebuild();

        return y;
    }

    TNode update(int val, int add, TNode x) {
        if (x == nullptr) x = new node(val);
        x->sz += add;
        x->sum += val * add;
        if (val == x->val) {
            x->cthis += add;
        }
        else if (val < x->val) {
            x->L = update(val, add, x->L);
            if (x->priority > x->L->priority) x = rotate_right(x);
        }
        else {
            x->R = update(val, add, x->R);
            if (x->priority > x->R->priority) x = rotate_left(x);
        }
        return x;
    }

    void add(int val) {
        root = update(val, 1, root);
    }

    void rem(int val) {
        root = update(val, -1, root);
    }

    long long sumKfirst(int cnt, TNode x) {
        if (!x) return 0;
        if (x->sz <= cnt) return x->sum;
        if (getsz(x->L) >= cnt) return sumKfirst(cnt, x->L);
        if (getsz(x->L) + x->cthis >= cnt) return getsum(x->L) + (cnt - getsz(x->L)) * x->val;
        return x->sum - getsum(x->R) + sumKfirst(cnt - getsz(x->L) - x->cthis, x->R);
    }

    long long sumKfirst(int cnt) {
        return sumKfirst(cnt, root);
    }
};

int n, isavail, avail[N], sz[N];
long long ans[N], in_sum[N];
long long dd[N][2];
vector<pii> S[N];
vector<int> List[N];
Treap cheapHihi[N];

void DFS(int u, int p, int max_sz) {
    avail[u] = isavail;
    if (u != p) --sz[u];

    int now_sz = min(sz[u], max_sz);
    vector<long long> Q;
    long long sum = 0, pre_sum = 0;
    REPD(i, S[u].size(), 0) {
        int c = S[u][i].first, v = S[u][i].second;
        if (v == p) continue;
        if (sz[v] > max_sz) {
            DFS(v, u, max_sz);
            sum += dd[v][1] + c;
            if (dd[v][0] - dd[v][1] - c < 0)
                Q.push_back(dd[v][0] - dd[v][1] - c);
        }
        else {
            in_sum[u] += c;
            swap(S[u][i], S[u].back());
            S[u].pop_back();

            cheapHihi[u].add(-c);
        }
    }

    sum += in_sum[u];
    if (max_sz) pre_sum = sum;
    else pre_sum = 1e15;

    for(int v : Q) cheapHihi[u].add(v);
    if (max_sz) pre_sum = sum + cheapHihi[u].sumKfirst(max_sz - 1);
    sum = sum + cheapHihi[u].sumKfirst(max_sz);
    for(int v : Q) cheapHihi[u].rem(v);

    dd[u][0] = pre_sum;
    dd[u][1] = sum;

    if (u != p) ++sz[u];
}

void sub34() {
    REP(i, 0, n) {
        sz[i] = S[i].size();
        REP(j, 0, sz[i]) List[j].push_back(i);
    }
    REP(i, 0, n) {
        ++isavail;
        for(int u : List[i]) if (isavail != avail[u]) {
            DFS(u, u, i);
            ans[i] += dd[u][1];
        }
    }
}

vector<long long> minimum_closure_costs(int N, vector<int> U, vector<int> V, vector<int> W) {
    n = N;
    REP(i, 0, n - 1) {
        S[U[i]].emplace_back(W[i], V[i]);
        S[V[i]].emplace_back(W[i], U[i]);
    }
    sub34();

    vector<long long> ret(n);
    REP(i, 0, n) ret[i] = ans[i];
    return ret;
}

Compilation message (stderr)

roads.cpp: In function 'void DFS(int, int, int)':
roads.cpp:106:9: warning: unused variable 'now_sz' [-Wunused-variable]
  106 |     int now_sz = min(sz[u], max_sz);
      |         ^~~~~~

• 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