Submission #412481

#	Time	Username	Problem	Language	Result	Execution time	Memory
412481		tatyam	Road Closures (APIO21_roads)	C++17	100 / 100	456 ms	43188 KiB

roads

#include <bits/stdc++.h>
using namespace std;
using ll = long long;

struct Edges{
    multiset<int> w, w2;
    ll sum = 0;
    void add(int x){
        sum += x;
        w.insert(x);
    }
    void erase(int x){
        if(auto p = w.find(x); p != w.end()){
            w.erase(p);
            sum -= x;
        }
        else{
            w2.erase(w2.find(x));
        }
    }
    pair<ll, vector<int>> topk_sum(int k, int deg){
        while(w2.size() && w.size() < k){
            auto p = prev(w2.end());
            sum += *p;
            w.insert(w.begin(), *p);
            w2.erase(p);
        }
        while(w.size() > k){
            auto p = w.begin();
            sum -= *p;
            w2.insert(w2.end(), *p);
            w.erase(p);
        }
        if(w.size() <= k - deg){
            return {sum, {}};
        }
        k = w.size() - (k - deg);
        vector<int> a(k);
        long b = sum;
        auto p = w.begin();
        for(int i = 0; i < k; i++) b -= a[i] = *p++;
        return {b, a};
    }
};
vector<ll> minimum_closure_costs(int N, vector<int> U, vector<int> V, vector<int> W){
    ll hidden = 0;
    vector<vector<pair<int, int>>> g(N), G(N);
    vector<Edges> e(N);
    for(int i = 0; i < N - 1; i++){
        const int a = U[i], b = V[i], c = W[i];
        G[a].emplace_back(b, c);
        G[b].emplace_back(a, c);
        e[a].add(c);
        e[b].add(c);
        hidden += c;
    }
    vector<int> idx(N);
    iota(idx.begin(), idx.end(), 0);
    sort(idx.begin(), idx.end(), [&](int x, int y){ return G[x].size() < G[y].size(); });
    set<int> roots;
    vector<int> uf(N, -1);
    auto root = [&](int x, auto root) -> int {
        return uf[x] < 0 ? x : uf[x] = root(uf[x], root);
    };
    auto unite = [&](int x, int y){
        x = root(x, root);
        y = root(y, root);
        assert(x != y);
        if(uf[x] > uf[y]) swap(x, y);
        uf[x] += uf[y];
        uf[y] = x;
        roots.erase(y);
    };
    vector<bool> awake(N);
    auto add_vertex = [&](int a){
        roots.insert(a);
        awake[a] = 1;
        for(auto& [b, c] : G[a]){
            e[b].erase(c);
            if(awake[b]){
                g[a].emplace_back(b, c);
                g[b].emplace_back(a, c);
                unite(a, b);
            }
            else{
                hidden -= c;
            }
        }
    };
    vector<ll> ans(N);
    for(int k = N; --k; ){
        auto dp = [&](int from, int a, int w, auto dp) -> pair<long, int> {
            const int deg = min<int>(g[a].size(), k);
            auto [sum, diff] = e[a].topk_sum(k, deg);
            if(deg == 0) return {sum, 0};
            for(auto [b, c] : g[a]) if(b != from){
                const auto [sum2, diff2] = dp(a, b, c, dp);
                sum += sum2;
                diff.push_back(diff2);
            }
            if(diff.size() < deg){
                for(int i : diff) sum += i;
                return {sum, w};
            }
            auto p = diff.begin() + (deg - 1);
            nth_element(diff.begin(), p, diff.end(), greater<>{});
            sum += accumulate(diff.begin(), p + 1, 0LL);
            return {sum, max(0, w - *p)};
        };
        while(idx.size() && G[idx.back()].size() == k){
            const int a = idx.back();
            idx.pop_back();
            add_vertex(a);
        }
        ans[k] = hidden;
        for(int a : roots){
            auto [sum, diff] = dp(-1, a, 0, dp);
            ans[k] += sum + diff;
        }
    }
    const ll sumW = accumulate(W.begin(), W.end(), 0LL);
    for(ll& i : ans) i = sumW - i;
    return ans;
}

Compilation message (stderr)

roads.cpp: In member function 'std::pair<long long int, std::vector<int> > Edges::topk_sum(int, int)':
roads.cpp:22:37: warning: comparison of integer expressions of different signedness: 'std::multiset<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   22 |         while(w2.size() && w.size() < k){
      |                            ~~~~~~~~~^~~
roads.cpp:28:24: warning: comparison of integer expressions of different signedness: 'std::multiset<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   28 |         while(w.size() > k){
      |               ~~~~~~~~~^~~
roads.cpp:34:21: warning: comparison of integer expressions of different signedness: 'std::multiset<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
   34 |         if(w.size() <= k - deg){
      |            ~~~~~~~~~^~~~~~~~~~
roads.cpp: In function 'std::vector<long long int> minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)':
roads.cpp:110:50: warning: comparison of integer expressions of different signedness: 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  110 |         while(idx.size() && G[idx.back()].size() == k){
      |                             ~~~~~~~~~~~~~~~~~~~~~^~~~
roads.cpp: In instantiation of 'minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)::<lambda(int, int, int, auto:24)> [with auto:24 = minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)::<lambda(int, int, int, auto:24)>]':
roads.cpp:117:47:   required from here
roads.cpp:101:28: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'const int' [-Wsign-compare]
  101 |             if(diff.size() < deg){
      |                ~~~~~~~~~~~~^~~~~

• 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