Submission #1223653

#	Time	Username	Problem	Language	Result	Execution time	Memory
1223653		PanosPask	Road Closures (APIO21_roads)	C++20	7 / 100	2096 ms	41924 KiB

roads

#include "roads.h"
#define pb push_back
#define mp make_pair

#include <vector>
#include <set>
#include <algorithm>

using namespace std;


typedef long long ll;
typedef pair<int, ll> pi;
const ll INF = 1e15;

struct SegTree {
	int size;
	vector<pi> tree;

	void init(int N) {
		size = 1;
		while (size < N) {
			size *= 2;
		}

		tree.assign(2 * size, mp(0, 0));
	}
};

int N;
vector<vector<pi>> adj_list;
vector<int> subtree;
vector<int> deg;

// dp1[node][k]: At most k roads open at the subtree of node, road from node to par is either open or closed, whichever is more optimal
vector<vector<ll>> dp1;

// dp2[node][k]: At most k roads open at the subtree of node, road from node to par is closed (IT IS counted in cost)
// dp2 goes up to deg(node), after that it's just the weight
vector<vector<ll>> dp2;

vector<vector<ll>> best1, best2;

void merge(vector<ll>& a, vector<ll>& b) {
	if (a.size() < b.size()) {
		swap(a, b);
	}

	for (int i = 0; i < b.size(); i++) {
		a[i] = a[i] + b[i];
	}
}

bool degcmp(pi a, pi b) {
	return deg[a.first] < deg[b.first];
}

inline void insert(multiset<ll>& s, ll& t, ll v, int top) {
	t += v;
	s.insert(v);	

	if (top < s.size()) {
		t -= *prev(s.end());
		s.erase(prev(s.end()));
	}
}

ll answer(int target, vector<ll>& diff, multiset<ll>& opt) {
	ll ans = INF;
	ll cur = 0;
	int c1 = 0;
	for (int i = 0; i < diff.size(); i++) {
		cur += diff[i];
		c1++;
	}

	auto it = opt.begin();
	int c2 = 0;
	while (c1 + c2 < target && it != opt.end()) {
		cur += *it;
		c2++;
		it++;
	}

	if (c1 + c2 == target) {
		ans = cur;
	}
	else {
		return INF;
	}

	for (int i = 0; i < diff.size(); i++) {
		cur -= diff[i];
		c1--;

		if (it == opt.end()) {
			return ans;
		}

		cur += *it;
		it++;
	}

	return ans;
}

// Calculate costs coming from children (both for d and d-1 children taken)
void calculate_best_cuts(int node) {
	multiset<ll> opt;
	ll sum = 0;

	best1[node].resize(deg[node] + 1);
	best2[node].resize(deg[node] + 1);

	sort(adj_list[node].begin(), adj_list[node].end(), degcmp);
	int start = 0;
	for (int d = 0; d <= deg[node]; d++) {
		vector<ll> diff;
		int p = start;
		while (p < adj_list[node].size()) {
			int u;
			ll w;
			tie(u, w) = adj_list[node][p];

			if (deg[u] == d) {
				start = p + 1;
				// Only need to keep the min deg[node] - d weights
				insert(opt, sum, w, deg[node] - d);
			}
			else {
				diff.pb(dp2[u][d] - dp1[u][d]);
			}
			p++;
		}
		sort(diff.begin(), diff.end());
		diff.resize(min((int)diff.size(), deg[node] - d));
		sort(diff.rbegin(), diff.rend());

		// Should have deg[node] - d elements (if at root, we have deg[node] - d + 1 elements)
		// If we are to open parent, we should look at closing deg[node] - d kids
		// If we are to close parent, we should look at closing deg[node] - d - 1 kids

		if (d == 0 && node != 0) {
			best1[node][0] = INF;
			best2[node][0] = 0;
			continue;
		}

		best1[node][d] = answer(deg[node] - d, diff, opt);

		if (d != deg[node]) {
			if (diff.size() == deg[node] - d) {
				diff.erase(diff.begin());
			}
			best2[node][d] = answer(deg[node] - d - 1, diff, opt);
		}
		else {
			best2[node][d] = 0;
		}

		if (node == 0) {
			best2[node][d] = INF;
		}
	}
}

void dfs(int node, int par, ll w) {
	deg[node] = adj_list[node].size();

	if (par != -1) {
		adj_list[node].erase(find(adj_list[node].begin(), adj_list[node].end(), mp(par, w)));
	}
		
	subtree[node] = 1;
	for (auto [neigh, w] : adj_list[node]) {
		dfs(neigh, node, w);
		subtree[node] += subtree[neigh];
	}
	calculate_best_cuts(node);

	vector<ll> dp;
	for (auto [neigh, w] : adj_list[node]) {
		merge(dp1[node], dp1[neigh]);
	}

	dp2[node].resize(deg[node] + 1);
	if (dp1[node].size() < deg[node] + 1) {
		dp1[node].resize(deg[node] + 1);
	}

	for (int d = 0; d <= deg[node]; d++) {
		dp2[node][d] = dp1[node][d] + best2[node][d] + w;
		dp1[node][d] = min(dp1[node][d] + best1[node][d], dp2[node][d]);
	}
}

vector<ll> minimum_closure_costs(int n, vector<int> U, vector<int> V,vector<int> W) {
	N = n;

	adj_list.resize(N);
	dp1.resize(N);
	dp2.resize(N);
	subtree.resize(N);
	deg.resize(N);
	best1.resize(N);
	best2.resize(N);

	for (int i = 0; i < N - 1; i++) {
		adj_list[U[i]].pb(mp(V[i], W[i]));
		adj_list[V[i]].pb(mp(U[i], W[i]));
	}

	dfs(0, -1, INF);
	while (dp1[0].size() < N) {
		dp1[0].pb(0);
	}

	return dp1[0];
}

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