Submission #437966

#	Time	Username	Problem	Language	Result	Execution time	Memory
437966		fleimgruber	Keys (IOI21_keys)	C++17	0 / 100	28 ms	35532 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;
	}
};

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 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
};
vector<Edge> out[MAX_N]; // outgoing edges. TODO optimise
int vis[MAX_N], par[MAX_N]; // vis_time, parent in DAG/tree
set<int> keys[MAX_N];
std::vector<int> in_comp[MAX_N]; // for each vertex the set of vertices
								 // it represents
// weakly is weakly connected components
UnionFind<int> compressed(MAX_N), weakly(MAX_N);

// merge i into j
// TODO smaller-to-larger
void merge(int i, int j) {
	compressed.connect(i, j);
	compressed.set_tag(j, j);
	in_comp[j].insert(in_comp[j].end(),
			in_comp[i].begin(), in_comp[i].end());
	in_comp[i].clear();
	for (int x : keys[i])
		keys[j].insert(x);
	keys[i].clear();
	// TODO no duplicates
	out[j].insert(out[j].end(), out[i].begin(), out[i].end());
	out[i].clear();
}

void solve() {
	bool trivial = false;
	for (int i = 0; i < n; i++) {
		bool found_out = false;
		for (auto& e : out[i])
			if (e.key == *keys[i].begin()) { // keys has size 1
				found_out = true;
				par[i] = e.to;
				weakly.connect(i, e.to);
				break;
			}
		if (!found_out) {
			trivial = true;
			ans[i] = 1;
		}
	}
	if (trivial)
		return;
	// otherwise each weakly connected component is a tree
	// with one extra edge. compress corresponding cycle to get a
	// directed forest
	// tags are the roots (where keys & edges are stored)
	for (int i = 0; i < n; i++) {
		compressed.set_tag(i, i);
		in_comp[i].push_back(i);
	}
	// roots, ordered by size
	struct PQKey {
		int vertex, size;
		bool operator<(const PQKey& o) const { return size > o.size; }
	};
	priority_queue<PQKey> roots;
	for (int i = 0; i < n; i++)
		if (vis[i] == 0) {
			int t = i + 1; // 1-based vis times
			int j = i;
			do {
				vis[j] = t;
				j = par[j];
			} while (vis[j] == 0);
			if (vis[j] == t) { // cycle!
				int start = j;
				for (j = par[j]; j != start; j = par[j])
					merge(j, start);
				par[start] = -1;
				roots.push({ start, compressed.size(start) });
			}
		}
	// != -1 means we already found an answer of size done_size
	// so stop pushing to pq, but collect remaining components
	int done_size = -1;
	while (!roots.empty()) {
		auto [r, _] = roots.top();
		roots.pop();
		// find any outgoing edge
		bool has_outgoing = false;
		for (auto& e : out[r]) {
			// TODO skip useless edges/duplicates
			if (compressed.connected(r, e.to) ||
					keys[r].find(e.key) == keys[r].end())
				continue;
			if (weakly.connected(r, e.to)) {
				// merge with self
				// walk up parent chain, until we're in root comp
				int j = compressed.tag(e.to);
				do {
					int jp = compressed.tag(par[j]);
					merge(j, jp);
					j = jp;
				} while (par[j] != -1);
				// we're still a root, but of larger size
				roots.push({ r, compressed.size(r) });				
			} else { // other component
				has_outgoing = true;
				// merge with other component, if not done
				if (done_size != -1)
					break;
				has_outgoing = true;
				par[r] = e.to;
				weakly.connect(r, e.to);
				// nothing new for priority queue, as we're
				// not a root anymore
				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]);
	int m = u.size();
	for (int i = 0; i < m; i++) {
		out[u[i]].push_back({ v[i], c[i] });
		out[v[i]].push_back({ u[i], c[i] });
	}
	ans.assign(n, 0);
	solve();
	return ans;
}

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