Submission #438080

#	Time	Username	Problem	Language	Result	Execution time	Memory
438080		fleimgruber	Keys (IOI21_keys)	C++17	67 / 100	2567 ms	111580 KiB

keys

// --- algolib/union-find.h ---
// union find, optionally with tags for each component. this is useful
// for example when compressing components. enable using template parameter
#include <algorithm>
#include <type_traits>
#include <numeric>
#include <vector>


template<typename T = void, typename enable = void>
class UnionFind {
	std::vector<int> p, size_;

public:
	UnionFind(int n = 0) : p(n), size_(n, 1) {
		std::iota(p.begin(), p.end(), 0);
	}

	int create() {
		int r = p.size();
		p.push_back(r);
		size_.push_back(1);
		return r;
	}

	int find(int i) {
		if (i == p[i])
			return i;
		return p[i] = find(p[i]);
	}
	int operator[](int i) { return find(i); }
	int size(int i) { return size_[find(i)]; }

	bool connected(int a, int b) { return find(a) == find(b); }

	bool connect(int a, int b) {
		a = find(a), b = find(b);
		if (a == b)
			return false;
		if (size_[a] > size_[b])
			std::swap(a, b);
		size_[b] += size_[a];
		p[a] = b;
		return true;
	}
};

// connect(a, b) maintains tag(b)
template<typename T>
class UnionFind<T, typename std::enable_if_t<!std::is_same_v<T, void>>>
	: public UnionFind<void> {
	std::vector<T> tags;

public:
	UnionFind(int n = 0) : UnionFind<void>(n), tags(n) { }

	void create() {
		UnionFind<void>::create();
		tags.emplace_back();
	}

	bool connect(int a, int b) {
		T old = tag(b);
		bool res = UnionFind<void>::connect(a, b);
		set_tag(b, old);
		return res;
	}

	void set_tag(int i, const T& t) { tags[find(i)] = t; }
	const T& tag(int i) { return tags[find(i)]; }
};
// ----------------
#include <bits/stdc++.h>

using namespace std;

const int MAX_N = 300005;

int n;
vector<int> ans;
struct Edge {
	int to, key; // key required
	bool operator<(const Edge& o) const {
		// cannot just use key, otherwise set breaks
		return tie(key, to) < tie(o.key, o.to);
	}
};
vector<int> out_can[MAX_N]; // only the to
set<Edge> out_blocked[MAX_N];
set<int> keys[MAX_N];
int par[MAX_N]; // parent in tree
std::vector<int> in_comp[MAX_N]; // for each vertex the set of vertices
								 // it represents (compressed)
UnionFind weakly(MAX_N); // weakly connected components
UnionFind<int> compressed(MAX_N);

// a into b
template<typename T>
void smaller_to_larger(vector<T>& a, vector<T>& b) {
	if (a.size() > b.size())
		a.swap(b);
	b.insert(b.end(), a.begin(), a.end());
	a.clear();
}
template<typename T>
void smaller_to_larger(set<T>& a, set<T>& b) {
	if (a.size() > b.size())
		a.swap(b);
	for (auto& x : a)
		b.insert(x);
	a.clear();
}

// merge i into j
void merge(int i, int j) {
	compressed.connect(i, j); // tag of j is maintained

	// easy ones
	smaller_to_larger(in_comp[i], in_comp[j]); // disjoint
	smaller_to_larger(out_can[i], out_can[j]);

	// TODO
	// should suffice up to subtask 4
	smaller_to_larger(out_blocked[i], out_blocked[j]);
	smaller_to_larger(keys[i], keys[j]);
	// promote blocked edges (for subtask 4 the keyset is small)
	for (int key : keys[j]) {
		auto it = out_blocked[j].lower_bound({ -1, key }); // v irrelevant
		while (it != out_blocked[j].end() && it->key == key) {
			out_can[j].push_back(it->to);
			it = out_blocked[j].erase(it);
		}
	}
	return;
}

void solve() {
	using pq_key = pair<int, int>; // size, vertex (roots order by size)
	priority_queue<pq_key, vector<pq_key>, greater<pq_key>> roots;
	// tags are the representatives (where keys & edges are stored)
	for (int i = 0; i < n; i++) {
		compressed.set_tag(i, i);
		in_comp[i].push_back(i);
		roots.push({ 1, i });
	}
	// != -1 means we already found an answer of size done_size
	// so stop pushing to pq, but collect remaining components
	int done_size = -1;
	ans.assign(n, 0);
	while (!roots.empty()) {
		auto [_, r] = roots.top();
		roots.pop();
		// find any outgoing edge
		bool has_outgoing = false;
		while (!out_can[r].empty()) {
			int to = out_can[r].back();
			out_can[r].pop_back();
			if (compressed.connected(r, to))
				continue;
			has_outgoing = true;
			if (done_size != -1) // we're larger than the answer
				break;
			if (weakly.connected(r, to)) {
				// merge with self. walk up parents 
				int j = compressed.tag(to);
				do {
					int jp = compressed.tag(par[j]);
					merge(j, jp);
					j = jp;
				} while (j != r);
				// we're still a root, but of larger size
				roots.push({ compressed.size(r), r });				
			} else { // merge with other component
				par[r] = to;
				weakly.connect(r, to);
				// we're no root any longer (no pq push)
			}
			break;
		}
		if (!has_outgoing && (done_size == -1 ||
					done_size == compressed.size(r))) {
			done_size = compressed.size(r);
			for (int x : in_comp[r])
				ans[x] = 1;
		}
	}
}

vector<int> find_reachable(vector<int> r, vector<int> u,
		vector<int> v, vector<int> c) {
	n = r.size();
	for (int i = 0; i < n; i++)
		keys[i].insert(r[i]);
	for (size_t i = 0; i < u.size(); i++)
		for (int x = 0; x < 2; x++) {
			if (r[u[i]] == c[i])
				out_can[u[i]].push_back(v[i]);
			else
				out_blocked[u[i]].insert({ v[i], c[i] });
			swap(u[i], v[i]);
		}
	solve();
	return ans;
}

• Subtask 1: 9 / 9
• Subtask 2: 11 / 11
• Subtask 3: 17 / 17
• Subtask 4: 30 / 30
• Subtask 5: 0 / 33