Submission #415335

#	Time	Username	Problem	Language	Result	Execution time	Memory
415335		ja_kingy	Road Closures (APIO21_roads)	C++14	100 / 100	385 ms	40076 KiB

roads

#include "roads.h"
#include <bits/stdc++.h>
using namespace std;

typedef long long ll;
typedef pair<ll,ll> pii;
const int mxN = 1e5;
ll sm, par[mxN], is_root[mxN], dp[mxN][2], done[mxN], par_cost[mxN];
list<int> order;
vector<pii> adj[mxN];
vector<ll> ans;

struct dp_val {
    ll v, dp0, dp1;
    bool operator < (dp_val b) const {
        ll da = dp1 - dp0, db = b.dp1 - b.dp0;
        if (da != db) return da > db;
        return v < b.v;
    }
};

set<dp_val> bestk[mxN], extra[mxN];

void dfs(int u, int p, int pc) {
    par[u] = p;
    par_cost[u] = pc;
    for (pii v: adj[u]) if (v.first != p) {
        dfs(v.first, u, v.second);
    }
    order.push_back(u);
}

void check_k(int u, int k) {
    if (bestk[u].size() < k && extra[u].size()) {
        auto fst = extra[u].begin();
        dp[u][0] -= fst->dp0;
        dp[u][0] += max(fst->dp1, fst->dp0);
        bestk[u].insert(*fst);
        extra[u].erase(fst);
    } 
    if (bestk[u].size() > k) {
        auto lst = bestk[u].end();
        lst--;
        dp[u][0] -= max(lst->dp1, lst->dp0);
        dp[u][0] += lst->dp0;
        extra[u].insert(*lst);
        bestk[u].erase(lst);
    }
    if (bestk[u].size() && extra[u].size()) {
        auto fst = extra[u].begin();
        auto lst = bestk[u].end();
        lst--;
        if (*fst < *lst) {
            dp[u][0] -= fst->dp0;
            dp[u][0] += max(fst->dp1, fst->dp0);
            bestk[u].insert(*fst);
            extra[u].erase(fst);
            dp[u][0] -= max(lst->dp1, lst->dp0);
            dp[u][0] += lst->dp0;
            extra[u].insert(*lst);
            bestk[u].erase(lst);
        }
    }
    if (bestk[u].size()) {
        auto lst = bestk[u].end();
        lst--;
        dp[u][1] = dp[u][0] - max(lst->dp1 - lst->dp0, 0ll) + par_cost[u];
    } else if (k) {
        dp[u][1] = par_cost[u];
    }
}

void erase_val(int u, int k, dp_val val) {
    auto v = bestk[u].find(val);
    if (v != bestk[u].end()) {
        dp[u][0] -= max(v->dp1, v->dp0);
        bestk[u].erase(v);
    } else {
        dp[u][0] -= val.dp0;
        extra[u].erase(val);
    }
    check_k(u, k);
}

void upd(int u, int k, dp_val val) {
    bestk[u].insert(val);
    dp[u][0] += max(val.dp1, val.dp0); 
    check_k(u, k);
}

vector<ll> minimum_closure_costs(int N, vector<int> U, vector<int> V, vector<int> W) {
    ans.resize(N);
    for (int i = 0; i < N-1; ++i) {
        adj[U[i]].push_back({V[i],W[i]});
        sm += W[i];
        adj[V[i]].push_back({U[i],W[i]});
    }
    dfs(0, -1, 0);
    is_root[0] = 1;
    for (int k = 0; k < N-1; ++k) {
        for (auto it = order.begin(); it != order.end();) {
            int u = *it;
            check_k(u, k);
            if (adj[u].size() == k) {
                done[u] = 1;
                for (pii v: adj[u]) {
                    if (done[v.first]) {
                        sm -= v.second;
                    } else {
                        upd(v.first, k, {u, 0, v.second});
                        if (par[v.first] == u) {
                            is_root[v.first] = 1;
                            ans[k] -= dp[v.first][0];
                        }
                    }
                }
                it = order.erase(it);
            } else if (is_root[u]) {
                ans[k] -= dp[u][0];
                it++;
            } else {
                upd(par[u], k, {u, dp[u][0], dp[u][1]});
                it++;
            }
        }
        for (auto it = order.rbegin(); it != order.rend(); it++) { 
            int u = *it;
            if (!is_root[u]) erase_val(par[u], k, {u, dp[u][0], dp[u][1]});
        }
        ans[k] += sm;
    }
    return ans;
}

Compilation message (stderr)

roads.cpp: In function 'void check_k(int, int)':
roads.cpp:34:25: warning: comparison of integer expressions of different signedness: 'std::set<dp_val>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   34 |     if (bestk[u].size() < k && extra[u].size()) {
      |         ~~~~~~~~~~~~~~~~^~~
roads.cpp:41:25: warning: comparison of integer expressions of different signedness: 'std::set<dp_val>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   41 |     if (bestk[u].size() > k) {
      |         ~~~~~~~~~~~~~~~~^~~
roads.cpp: In function 'std::vector<long long int> minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)':
roads.cpp:104:31: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  104 |             if (adj[u].size() == k) {
      |                 ~~~~~~~~~~~~~~^~~~

• 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