Submission #528507

#	Time	Username	Problem	Language	Result	Execution time	Memory
528507		cologne	Road Closures (APIO21_roads)	C++17	100 / 100	155 ms	44632 KiB

roads

#include "roads.h"

#include <algorithm>
#include <functional>
#include <map>
#include <utility>
#include <vector>
using namespace std;

struct Choicer
{
    map<long long, int> M;
    map<long long, int>::iterator ptr;
    long long left_sum;
    int left_count;
    long long total_sum;

    Choicer() : M(), ptr(M.end()), left_sum(0), left_count(0), total_sum(0)
    {
    }

    void add(long long pick, long long skip, int delta)
    {
        pick = min(pick, skip);
        total_sum += skip * delta;
        long long save_cost = pick - skip;
        if (ptr == M.end() || save_cost < ptr->first)
        {
            left_sum += save_cost * delta;
            left_count += delta;
        }

        auto it = M.insert({save_cost, 0LL}).first;
        it->second += delta;

        if (it->second == 0)
        {
            if (it == ptr)
                ptr = next(ptr);
            M.erase(it);
        }
    }

    long long sum(int limit)
    {
        if (limit == 0)
            return total_sum;

        while (left_count >= limit)
        {
            ptr = prev(ptr);
            left_count -= ptr->second;
            left_sum -= ptr->first * ptr->second;
        }
        while (ptr != M.end() && left_count + ptr->second < limit)
        {
            left_count += ptr->second;
            left_sum += ptr->first * ptr->second;
            ptr = next(ptr);
        }
        if (ptr == M.end())
            return total_sum + left_sum;
        return total_sum + left_sum + (limit - left_count) * ptr->first;
    }
};

vector<long long> minimum_closure_costs(int N, vector<int> U, vector<int> V, vector<int> W)
{
    vector<vector<pair<int, int>>> adj(N);
    vector<long long> answer(N);
    for (int i = 0; i < N - 1; i++)
    {
        adj[U[i]].emplace_back(V[i], W[i]);
        adj[V[i]].emplace_back(U[i], W[i]);
        answer[0] += W[i];
    }

    vector<vector<int>> targets(N);
    vector<vector<int>> exactDegree(N);
    for (int i = 0; i < N; i++)
    {
        sort(adj[i].begin(), adj[i].end(), [&](pair<int, int> a, pair<int, int> b)
             { return adj[a.first].size() > adj[b.first].size(); });
        for (int j = 0; j < (int)adj[i].size(); j++)
            targets[j].push_back(i);
        exactDegree[adj[i].size()].push_back(i);
    }

    vector<bool> vis(N, true);
    vector<Choicer> choice(N);
    for (int k = 1; k < N; k++)
    {

        function<pair<long long, long long>(int, int)> dfs = [&](int a, int p)
        {
            vis[a] = true;

            vector<pair<long long, long long>> costs;
            for (auto [x, w] : adj[a])
            {
                if ((int)adj[x].size() <= k)
                    break;
                if (x == p)
                    continue;
                auto [lt, le] = dfs(x, a);
                costs.emplace_back(lt, le + w);
            }

            for (auto [pick, skip] : costs)
                choice[a].add(pick, skip, +1);

            long long lt = choice[a].sum(k - 1), le = choice[a].sum(k);

            for (auto [pick, skip] : costs)
                choice[a].add(pick, skip, -1);

            return make_pair(lt, le);
        };

        for (int u : exactDegree[k])
            for (auto [v, w] : adj[u])
                if (adj[v].size() > adj[u].size())
                    choice[v].add(0, w, +1);

        for (int u : targets[k])
            vis[u] = false;
        for (int u : targets[k])
            if (!vis[u])
                answer[k] += dfs(u, -1).second;
    }

    return answer;
}

• 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