Submission #955223

#	Time	Username	Problem	Language	Result	Execution time	Memory
955223		alontanay	Road Closures (APIO21_roads)	C++14	21 / 100	2085 ms	51364 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();
    }
};

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(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();
    }
};

vector<pi> nei[MAX_N];
vector<pi> bigNei[MAX_N];
int vis[MAX_N];
SlowKSum smalls[MAX_N];

pl dfs(int node, int K, int par = -1) {
    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) {
    vector<ll> ans(N);
    vector<int> big(N);

    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);
    }
    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);
            big[node] = true;
            bigs.push_back(node);
            for (pi edge : nei[node]) {
                int ne = edge.f, weight = edge.s;
                if (!big[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:49:24: warning: comparison of integer expressions of different signedness: 'std::multiset<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   49 |         if(vals.size() < _k) { return INF_LL; }
      |            ~~~~~~~~~~~~^~~~
roads.cpp: In function 'pl dfs(int, int, int)':
roads.cpp:162:24: warning: comparison of integer expressions of different signedness: 'std::vector<long long int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  162 |     if(bigCosts.size() > dif) {
      |        ~~~~~~~~~~~~~~~~^~~~~
roads.cpp:179:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  179 |     for(int i = 0; i < bigCosts.size(); i ++) {
      |                    ~~^~~~~~~~~~~~~~~~~

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