Submission #725102

#	Time	Username	Problem	Language	Result	Execution time	Memory
725102		faremy	Road Closures (APIO21_roads)	C++14	100 / 100	430 ms	44004 KiB

roads

#include <iostream>
#include <random>

using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;

//Wrapper for generating random integers from [0, N[
class RandInt
{
        static std::random_device rd;
        static std::mt19937 gen;
        std::uniform_int_distribution<ll> distrib;

public:
        RandInt(ll N) : distrib(0, N - 1) {}
        ll operator ()() { return distrib(gen); }
};

std::random_device RandInt::rd;
std::mt19937 RandInt::gen(rd());

//End of template

#include "roads.h"

#include <vector>
#include <utility>
#include <unordered_map>
#include <queue>
#include <algorithm>

//Treap
struct TreapNode
{
	TreapNode(ll v) : value(v), priority(gen_prio()), sum(v) {}
	~TreapNode() { delete left; delete right; }

	ll value, priority;
	int size = 1;
	ll sum;
	TreapNode *left = nullptr, *right = nullptr;

private:
	static RandInt gen_prio;
};

const ll INFTY = 1LL << 62;
RandInt TreapNode::gen_prio(INFTY);

using Treap = TreapNode *;
using Treapair = std::pair<Treap, Treap>;

int size(Treap t)
{
	if (t)
		return t->size;
	return 0;
}

ll sum(Treap t)
{
	if (t)
		return t->sum;
	return 0;
}

Treapair cut(Treap t, ll v)
{
	if (t == nullptr)
		return { nullptr, nullptr };
	if (t->value < v)
	{
		Treapair sub = cut(t->right, v);
		t->right = sub.first;
		t->size -= size(sub.second);
		t->sum -= sum(sub.second);
		return { t, sub.second };
	}
	else
	{
		Treapair sub = cut(t->left, v);
		t->left = sub.second;
		t->size -= size(sub.first);
		t->sum -= sum(sub.first);
		return { sub.first, t };
	}
}

Treap merge(Treap l, Treap r)
{
	if (l == nullptr)
		return r;
	if (r == nullptr)
		return l;
	if (l->priority > r->priority)
	{
		l->size += size(r);
		l->sum += sum(r);
		l->right = merge(l->right, r);
		return l;
	}
	else
	{
		r->size += size(l);
		r->sum += sum(l);
		r->left = merge(l, r->left);
		return r;
	}
}

void insert(Treap &t, ll v)
{
	Treapair split = cut(t, v);
	t = merge(merge(split.first, new TreapNode(v)), split.second);
}

int cnt_range(Treap &t, ll a, ll b)
{
	Treapair low = cut(t, a);
	Treapair high = cut(low.second, b);
	int res = size(high.first);
	t = merge(low.first, merge(high.first, high.second));
	return res;
}

ll sum_k(Treap &t, int k)
{
	if (k <= 0)
		return 0;
	if (t == nullptr)
		return INFTY;
	
	if (k > size(t->left))
		return sum(t->left) + t->value + sum_k(t->right, k - size(t->left) - 1);
	return sum_k(t->left, k);
}

ll sum_k_range(Treap &t, ll a, ll b, int k)
{
	Treapair low = cut(t, a);
	Treapair high = cut(low.second, b);
	ll res = sum_k(high.first, k);
	t = merge(low.first, merge(high.first, high.second));
	return res;
}

//Main code
const int MAXN = 1e5;

std::unordered_map<int, ll> tree[MAXN];
Treap dead_neighbors[MAXN];
int degree[MAXN];
int deg_order[MAXN];

bool seen[MAXN];
ll min_cuts[MAXN][2];

ll cut_children(int node, int parent, int maxdeg)
{
	int remain = degree[node] - maxdeg;
	std::priority_queue<ll> add_cost;
	ll res = 0;

	for (auto [child, _] : tree[node])
	{
		if (child != parent)
		{
			if (min_cuts[child][1] > min_cuts[child][0])
			{
				res += min_cuts[child][0];
				add_cost.emplace(min_cuts[child][0] - min_cuts[child][1]);
			}
			else
			{
				res += min_cuts[child][1];
				remain--;
			}
		}
	}

	ll previous_alive = 0;
	while (remain > 0 && !add_cost.empty())
	{
		ll smallest_alive = -add_cost.top();
		add_cost.pop();

		int cnt_dead = cnt_range(dead_neighbors[node], previous_alive, smallest_alive);
		if (cnt_dead < remain)
		{
			res += sum_k_range(dead_neighbors[node], previous_alive, smallest_alive, cnt_dead) + smallest_alive;
			remain -= cnt_dead + 1;
		}
		else
		{
			res += sum_k_range(dead_neighbors[node], previous_alive, smallest_alive, remain);
			remain = 0;
		}

		previous_alive = smallest_alive;
	}

	res += sum_k_range(dead_neighbors[node], previous_alive, INFTY, remain);
	return res;
}

void compute_cuts(int node, int parent, int maxdeg)
{
	seen[node] = true;
	for (auto [child, _] : tree[node])
		if (child != parent)
			compute_cuts(child, node, maxdeg);

	min_cuts[node][0] = cut_children(node, parent, maxdeg);
	if (parent != -1)
		min_cuts[node][1] = cut_children(node, parent, maxdeg + 1) + tree[node][parent];
}

std::vector<ll> minimum_closure_costs(int N, std::vector<int> U, std::vector<int> V, std::vector<int> W)
{
	for (int iEdge = 0; iEdge < N - 1; iEdge++)
	{
		tree[U[iEdge]][V[iEdge]] = W[iEdge];
		tree[V[iEdge]][U[iEdge]] = W[iEdge];
	}
	
	for (int iNode = 0; iNode < N; iNode++)
	{
		dead_neighbors[iNode] = nullptr;
		degree[iNode] = tree[iNode].size();
		deg_order[iNode] = iNode;
		seen[iNode] = false;
	}

	std::sort(deg_order, deg_order + N, [](int u, int v) { return degree[u] < degree[v]; });
	int first_alive = 0;

	std::vector<ll> ans(N, 0);
	for (int maxdeg = 0; maxdeg < N; maxdeg++)
	{
		while (first_alive < N && degree[deg_order[first_alive]] <= maxdeg)
		{
			for (auto [neighbor, weight] : tree[deg_order[first_alive]])
			{
				tree[neighbor].erase(tree[neighbor].find(deg_order[first_alive]));
				insert(dead_neighbors[neighbor], weight);
			}
			first_alive++;
		}

		for (int iNode = first_alive; iNode < N; iNode++)
			if (!seen[deg_order[iNode]])
			{
				compute_cuts(deg_order[iNode], -1, maxdeg);
				ans[maxdeg] += min_cuts[deg_order[iNode]][0];
			}

		for (int iNode = first_alive; iNode < N; iNode++)
			seen[deg_order[iNode]] = false;
	}

	for (int iNode = 0; iNode < N; iNode++)
		delete dead_neighbors[iNode];
	return ans;
}

Compilation message (stderr)

roads.cpp: In function 'll cut_children(int, int, int)':
roads.cpp:165:12: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  165 |  for (auto [child, _] : tree[node])
      |            ^
roads.cpp: In function 'void compute_cuts(int, int, int)':
roads.cpp:210:12: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  210 |  for (auto [child, _] : tree[node])
      |            ^
roads.cpp: In function 'std::vector<long long int> minimum_closure_costs(int, std::vector<int>, std::vector<int>, std::vector<int>)':
roads.cpp:243:14: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  243 |    for (auto [neighbor, weight] : tree[deg_order[first_alive]])
      |              ^

• 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