Submission #955876

#	Time	Username	Problem	Language	Result	Execution time	Memory
955876		alontanay	Road Closures (APIO21_roads)	C++14	31 / 100	2040 ms	53492 KiB

roads

#include "roads.h"

#include <bits/stdc++.h>
// #define DEBUGGING


#ifdef DEBUGGING
#define dout(msg) cout << msg
#else
#define dout(msg) 
#endif

#define f first
#define s second

using namespace std;
using ll = long long;
using pi = pair<int,int>;
using pl = pair<ll,ll>;

const int MAX_N = 1e5 + 5;

const ll INF_LL = 1e18;
const int INF = 2e9;

struct SlowKSum {
    int _k = 0;
    multiset<int> _vals;
    int getK() { return _k; }
    void incK() {
        _k++;
    }
    void decK() {
        _k--;
    }
    void setK(int k) {
        while(_k < k) {
            incK();
        }
        while(_k > k) {
            decK();
        }
    }
    void insert(int x) {
        _vals.insert(x);
    }
    void erase(int x) {
        _vals.erase(_vals.find(x));
    }
    ll getSum() {
        if(_vals.size() < _k) { return INF_LL; }
        ll res = 0;
        auto it = _vals.begin();
        for(int i = 0; i < _k; i ++) {
            res += *it;
            it++;
        }
        return res;
    }
    ll getKSum(int k) {
        setK(k);
        return getSum();
    }
    void pre() {
    }
    void post() {
    }
};

struct KSum {
   private:
    int _k = 0;
    ll _sum = 0;
    int _overflowCount = 0;
    int _opCount = 0;
    set<pi> _vals;
    set<pi>::iterator sep = _vals.end();

    void incSep() {
        if (sep == _vals.end()) {
            _overflowCount++;
            return;
        }
        _sum += sep->f;
        sep++;
    };
    void decSep() {
        if (_overflowCount) {
            _overflowCount--;
            return;
        }
        assert(sep != _vals.begin());
        sep--;
        _sum -= sep->f;
    }
    int sepVal() {
        if(_overflowCount || sep == _vals.end()) { return INF; }
        return sep->f;
    }
   public:
    int getK() { return _k; }
    void incK() {
        _k++;
        incSep();
    }
    void decK() {
        _k--;
        decSep();
    }
    void setK(int k) {
        while(_k < k) {
            incK();
        }
        while(_k > k) {
            decK();
        }
    }
    void insert(int x) {
        _vals.insert({x,_opCount++});
        if(x < sepVal()) {
            _sum += x;
            decSep();
        }
    }
    void erase(int x) {
        if(x <= sepVal()) {
            _sum -= x;
            incSep();
        }
        _vals.erase(_vals.lower_bound({x,-INF}));
    }
    ll getSum() {
        return _overflowCount ? INF_LL : _sum;
    }
    ll getKSum(int k) {
        setK(k);
        return getSum();
    }
    void pre() {
        _k = INF;
        _overflowCount = INF;
    }
    void post() {
        _k = 0;
        _overflowCount = 0;
        _sum = 0;
        sep = _vals.begin();
    }
};

template<const int BOUND>
struct BoundKSum {
   private:
    int _k = 0;
    vector<ll> hist = vector<ll>(BOUND);
   public:
    int getK() { return _k; }
    void incK() {
        _k++;
    }
    void decK() {
        _k--;
    }
    void setK(int k) {
        while(_k < k) {
            incK();
        }
        while(_k > k) {
            decK();
        }
    }
    void insert(int x) {
        hist[x]++;
    }
    void erase(int x) {
        hist[x]--;
    }
    ll getSum() {
        ll lft = _k;
        ll res = 0;
        for(ll i = 0; i < BOUND; i ++) {
            ll cnt = min(lft,hist[i]);
            lft -= cnt;
            res += cnt * i;
        }
        return lft ? INF_LL : res;
    }
    ll getKSum(int k) {
        setK(k);
        return getSum();
    }
    void pre() {
    }
    void post() {
    }
};


vector<pi> nei[MAX_N];
int isBig[MAX_N];
vector<pi> bigNei[MAX_N];
int vis[MAX_N];
KSum smalls[MAX_N];
int _N;

int _dfs_count = 0;

pl dfs(int node, int K, int par = -1) {
    _dfs_count++;
    assert(_dfs_count < 4*_N);
    // assert(_dfs_count < 4*_N);
    assert(nei[node].size() > K);
    vis[node] = true;

    int deg = bigNei[node].size();
    int dif = nei[node].size()-K;
    dout(node << ' ' << K << ' ' << par << endl);
    assert(dif > 0);
    ll total = 0;
    vector<ll> bigCosts;
    bigCosts.reserve(deg - (par != -1));
    ll parentWeight = INF_LL;
    for(pi edge : bigNei[node]) {
        int ne = edge.f, w = edge.s;
        if(ne == par) {
            parentWeight = w;
            continue;
        }
        pl parts = dfs(ne, K, node);
        total += parts.f;
        bigCosts.push_back(parts.s);
    }
    sort(bigCosts.begin(),bigCosts.end());
    if(bigCosts.size() > dif) {
        bigCosts.resize(dif);
    }
    dout(bigCosts.size() << endl);
    vector<ll> smallCosts(dif+1);
    for(int i = dif; i >= 0; i --) {
        smallCosts[i] = smalls[node].getKSum(dif-i);
    }
    dout("NODE " << node << endl);
    dout("DIF " << dif  << endl);
    for(int i = 0; i <= dif; i ++) {
        dout(smallCosts[i] << ' ');
    }
    dout(endl);
    ll best = smallCosts[0];
    ll bestWithParent = smallCosts[1] + parentWeight;
    ll preSum = 0;
    for(int i = 0; i < bigCosts.size(); i ++) {
        dout("cost " << bigCosts[i] << endl);
        preSum += bigCosts[i];
        best = min(best, preSum + smallCosts[i+1]);
        if(i < dif-1) { 
            bestWithParent = min(bestWithParent, preSum + parentWeight + smallCosts[i+2]);
        }
    }
    best = min(best,bestWithParent);
    dout("BEST " << best << endl);
    return {best+total,bestWithParent-best};
}

vector<ll> minimum_closure_costs(int N, vector<int> U, vector<int> V,
                                 vector<int> W) {
    _N = N;
    vector<ll> ans(N);

    for(int i = 0; i < N; i ++) {
        smalls[i].pre();
    }
    for (int i = 0; i < N - 1; i++) {
        int a = U[i], b = V[i], h = W[i];
        nei[a].push_back({b,h});
        nei[b].push_back({a,h});
        smalls[a].insert(h);
        smalls[b].insert(h);
    }
    for(int i = 0; i < N; i ++) {
        smalls[i].post();
    }
    vector<vector<int>> degHist(N);
    for (int i = 0; i < N; i++) {
        int d = nei[i].size();
        degHist[d].push_back(i);
        bigNei[i].reserve(d);
    }
    vector<int> bigs;
    bigs.reserve(N);
    for (int K = N - 1; K >= 0; K--) {
        dout("~~~~~~~~ " << K << endl);
        ll res = 0;
        for(int big : bigs) {
            if(!vis[big]) {
                res += dfs(big,K).f;
            }
        }
        ans[K] = res;
        for(int big : bigs) {
            vis[big] = false;
        }
        for (int node : degHist[K]) {
            dout(" + " << node << endl);
            isBig[node] = true;
            bigs.push_back(node);
            for (pi edge : nei[node]) {
                int ne = edge.f, weight = edge.s;
                if (!isBig[ne]) {
                    continue;
                }
                bigNei[node].push_back({ne,weight});
                bigNei[ne].push_back({node,weight});
                smalls[node].erase(weight);
                smalls[ne].erase(weight);
            }
        }
    }
    return ans;
}

Compilation message (stderr)

roads.cpp: In member function 'll SlowKSum::getSum()':
roads.cpp:51:25: warning: comparison of integer expressions of different signedness: 'std::multiset<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   51 |         if(_vals.size() < _k) { return INF_LL; }
      |            ~~~~~~~~~~~~~^~~~
In file included from /usr/include/c++/10/cassert:44,
                 from /usr/include/x86_64-linux-gnu/c++/10/bits/stdc++.h:33,
                 from roads.cpp:3:
roads.cpp: In function 'pl dfs(int, int, int)':
roads.cpp:212:29: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  212 |     assert(nei[node].size() > K);
      |            ~~~~~~~~~~~~~~~~~^~~
roads.cpp:234:24: warning: comparison of integer expressions of different signedness: 'std::vector<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  234 |     if(bigCosts.size() > dif) {
      |        ~~~~~~~~~~~~~~~~^~~~~
roads.cpp:251:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  251 |     for(int i = 0; i < bigCosts.size(); i ++) {
      |                    ~~^~~~~~~~~~~~~~~~~

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