Submission #708421

#	Time	Username	Problem	Language	Result	Execution time	Memory
708421		600Mihnea	Road Closures (APIO21_roads)	C++17	100 / 100	289 ms	35120 KiB

roads

#include "roads.h"

#include <cmath>
#include <functional>
#include <fstream>
#include <iostream>
#include <vector>
#include <algorithm>
#include <string>
#include <set>
#include <map>
#include <list>
#include <time.h>
#include <math.h>
#include <random>
#include <deque>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <iomanip>
#include <cassert>
#include <bitset>
#include <sstream>
#include <chrono>
#include <cstring>
#include <numeric>

using namespace std;

typedef long long ll;

#define all(x) x.begin(), x.end()
#define f first
#define s second
#define pb push_back
#define pii pair<int, int>
#define ve vector
#define vei vector<int> 
#define vell vector<ll> 
#define sz(x) ((int)x.size())
#define each(it, x) for (auto &it : x)
#define debug(x) cout << #x << " = " << x << "\n";
#define FU function
#define VO void

vell minimum_closure_costs(int N, vei U, vei V, vei W) {
    int n = N;
    assert(sz(U) == n - 1);
    assert(sz(V) == n - 1);
    assert(sz(W) == n - 1);
#ifdef ONPC
    if (0) {
        debug(n);
        each(it, U) cout << it << " "; cout << "\n";
        each(it, V) cout << it << " "; cout << "\n";
        each(it, W) cout << it << " "; cout << "\n";
    }
#endif
    ve<ve<pii>> g(n);
    for (int i = 0; i < n - 1; i++)
        g[U[i]].pb({ V[i], W[i] }),
        g[V[i]].pb({ U[i], W[i] });

    vei ord, deg(n, 0), pap(n, -1), upp(n, 0);
    for (int i = 0; i < n; i++) deg[i] = sz(g[i]);
    vector<bool> vis(n, 0);

    ord.push_back(0);
    vis[0] = 1;
    for (int i = 0; i < (int)ord.size(); i++)
    {
        int a = ord[i];
        ve<pii> kids;
        each(it, g[a])
        {
            int b = it.f, c = it.s;
            if (vis[b]) continue;
            pap[b] = a;
            upp[b] = c;
            vis[b] = 1;
            ord.pb(b);
            kids.pb({ b, c });
        }
        g[a] = kids;
    }
    assert((int)ord.size() == n);
    reverse(all(ord));
    vei loc(n);
    for (int i = 0; i < n; i++) loc[ord[i]] = i;

    for (int i = 0; i < n; i++)
        sort(g[i].begin(), g[i].end(), [&](pair<int, int> i, pair<int, int> j) {return deg[i.first] > deg[j.first]; });

    vector<vector<int>> devine(n);
    for (int i = 0; i < n; i++)
    {
        devine[deg[i]].pb(i);
    }

    vell sol(n, 0);
    vei verts;

    vell dp0(n, 0), dp1(n, 0);

    ll smsz = 0;
    vector<ll> sui(n, 0);
    vector<multiset<int>> gustav1(n), gustav2(n);
    vector<ll> s1(n, 0);
    for (int a = 0; a < n; a++) each(it, g[a]) sui[a] += it.second, gustav2[a].insert(-it.second);

    auto mv12 = [&](int a)
    {
        assert(!gustav1[a].empty());
        auto it = gustav1[a].end();
        it--;
        int x = *it;
        gustav1[a].erase(it);
        gustav2[a].insert(x);
        s1[a] -= x;
    };

    auto mv21 = [&](int a)
    {
        assert(!gustav2[a].empty());
        auto it = gustav2[a].begin();
        int x = *it;
        gustav2[a].erase(it);
        gustav1[a].insert(x);
        s1[a] += x;
    };


    for (int k = n - 1; k >= 0; k--)
    {
        each(x, devine[k])
        {
            verts.pb(x);
            if (pap[x] != -1)
            {
                if (gustav1[pap[x]].find(-upp[x]) != gustav1[pap[x]].end())
                {
                    s1[pap[x]] -= (-upp[x]);
                    gustav1[pap[x]].erase(gustav1[pap[x]].find(-upp[x]));
                }
                else
                    if (gustav2[pap[x]].find(-upp[x]) != gustav2[pap[x]].end())
                    {
                        gustav2[pap[x]].erase(gustav2[pap[x]].find(-upp[x]));
                    }
                    else assert(0);
                sui[pap[x]] -= upp[x];
            }
        }

        smsz += sz(verts);
        sort(all(verts), [&](int a, int b) {return loc[a] < loc[b]; });

        // 0 => are gradul <= k + 0
        // 1 => are gradul <= k + 1
        ll print = 0;
        each(a, verts)
        {
            assert(deg[a] >= k);
            each(rem, vei({ k - 1, k }))
            {
                // sterg muchia => dp1 + cost
                // pastrez muchia => dp0

                int total = sz(g[a]);
                int pastrez = min(total, rem);

                pastrez = max(pastrez, 0);

                ll su = sui[a];
                each(it, g[a]) if (deg[it.first] >= k) su += dp1[it.first] + it.second; else break;
                vector<ll> difs;
                each(it, g[a]) if (deg[it.first] >= k) difs.pb(min(0LL, dp0[it.first] - dp1[it.first] - it.second)); else break;

                auto iter = gustav1[a].begin();
                sort(difs.begin(), difs.end());

                int po = 0;

                ll best = (ll)1e18;

                ll cur = 0;

                assert((int)difs.size() + (int)gustav1[a].size() + (int)gustav2[a].size() >= pastrez);

                for (int take = 0; take <= pastrez; take++)
                {
                    if (take > (int)difs.size()) break;
                    if (take)
                    {
                        cur += difs[take - 1];
                    }
                    int from_gustav = pastrez - take;
                    if (from_gustav <= (int)gustav1[a].size() + (int)gustav2[a].size())
                    {
                        while ((int)gustav1[a].size() < from_gustav) mv21(a);
                        while ((int)gustav1[a].size() > from_gustav) mv12(a);
                        assert((int)gustav1[a].size() == from_gustav);

                        best = min(best, cur + s1[a]);
                    }
                }
                su += best;
                if (rem == k - 1) dp0[a] = su; else dp1[a] = su;
            }
            if (a == 0)
            {
                print += dp1[a];
            }
            else
            {
                if (deg[pap[a]] < k)
                {
                    print += min(dp0[a], dp1[a] + upp[a]);
                }
            }
            assert(dp0[a] >= dp1[a]);
        }
        sol[k] = print;
    }
    assert(smsz <= 3 * n);
    return sol;
}

Compilation message (stderr)

roads.cpp: In function 'std::vector<long long int> minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)':
roads.cpp:179:22: warning: variable 'iter' set but not used [-Wunused-but-set-variable]
  179 |                 auto iter = gustav1[a].begin();
      |                      ^~~~
roads.cpp:182:21: warning: unused variable 'po' [-Wunused-variable]
  182 |                 int po = 0;
      |                     ^~

• 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