Keys
Given n rooms and m bidirectional corridors. Each room i contains a key of type r[i]. Each corridor j connects rooms u[j] and v[j] and requires a key of type c[j] to traverse. Starting from room i, you pick up the key...
Problem Statement
Rendered from the "Problem Summary" section in the LaTeX write-up.
Given $n$ rooms and $m$ bidirectional corridors. Each room $i$ contains a key of type $r[i]$. Each corridor $j$ connects rooms $u[j]$ and $v[j]$ and requires a key of type $c[j]$ to traverse.
Starting from room $i$, you pick up the key in room $i$, then traverse any corridor for which you have the required key, pick up keys in visited rooms, and continue. Let $p[i]$ be the number of rooms reachable from room $i$.
Find all rooms $i$ that have the minimum $p[i]$.
Editorial
The solution write-up is rendered from the LaTeX source, with equations kept live through MathJax.
Solution Approach
BFS/DFS with Union-Find
For each starting room, we could run a BFS collecting keys and exploring new corridors. But doing this independently for each room is $O(n(n + m))$.
Efficient Approach: Merging Components
The key insight is that if two rooms are mutually reachable, they reach the same set of rooms. We use a Union-Find structure:
For each room $i$, start a BFS/DFS exploration.
Maintain the set of collected keys and frontier corridors.
When a new key is collected, check all corridors requiring that key type.
If exploration from room $i$ reaches room $j$, and we previously explored from $j$ and found its reachable set, merge them.
Optimized Algorithm
Process rooms in a smart order. Start BFS from each room. If during BFS from room $i$, we reach a room $j$ that was already fully explored (its reachable set is known and smaller), then $p[i] \ge p[j]$. Conversely, if we reach a room $j$ that is currently being explored, merge the explorations.
The algorithm uses a ``small-to-large'' merging strategy:
Start BFS from room 0. Maintain key set $K$, visited rooms $V$, and pending corridors $E$.
When a new key of type $t$ is added to $K$, add all corridors of type $t$ incident to visited rooms to the frontier.
If BFS terminates (no more reachable rooms), $p[0] = |V|$.
If during BFS we enter a room already in another component, merge components.
C++ Solution
#include <bits/stdc++.h>
using namespace std;
vector<int> find_reachable(vector<int> r, vector<int> u,
vector<int> v, vector<int> c) {
int n = r.size(), m = u.size();
// Group corridors by key type
map<int, vector<int>> corridors_by_type; // type -> list of corridor indices
for (int j = 0; j < m; j++) {
corridors_by_type[c[j]].push_back(j);
}
// Adjacency list: for each room, corridors incident to it
vector<vector<int>> adj(n);
for (int j = 0; j < m; j++) {
adj[u[j]].push_back(j);
adj[v[j]].push_back(j);
}
vector<int> reach(n, 0);
vector<int> p(n, n); // reachable count
// For each room, BFS
// Optimization: mark rooms whose reachable set is already determined
vector<bool> done(n, false);
for (int start = 0; start < n; start++) {
if (done[start]) continue;
set<int> keys;
set<int> visited;
queue<int> bfs;
queue<int> pending; // corridors with matching key waiting to be processed
auto visit_room = [&](int room) {
if (visited.count(room)) return;
visited.insert(room);
int key = r[room];
if (!keys.count(key)) {
keys.insert(key);
// Activate all corridors of this key type that connect to visited rooms
if (corridors_by_type.count(key)) {
for (int j : corridors_by_type[key]) {
if (visited.count(u[j]) || visited.count(v[j]))
pending.push(j);
}
}
}
// Add corridors from this room that match collected keys
for (int j : adj[room]) {
if (keys.count(c[j])) {
pending.push(j);
}
}
};
visit_room(start);
while (!pending.empty()) {
int j = pending.front(); pending.pop();
int a = u[j], b = v[j];
if (!visited.count(a)) {
visit_room(a);
}
if (!visited.count(b)) {
visit_room(b);
}
}
int cnt = visited.size();
for (int room : visited) {
p[room] = min(p[room], cnt);
// If this room reaches 'cnt' rooms, and we explored from it,
// any room in visited that starts exploration will also reach at least cnt.
// But they might reach more if they haven't been explored yet.
}
}
// Actually, the above is incorrect for rooms not equal to start.
// Correct approach: p[i] = size of reachable set from room i.
// We need to compute this for every room.
// The BFS above from start gives p[start] = visited.size().
// For other rooms in visited, they might reach MORE rooms (since they have different keys).
// Recompute properly: BFS from each room independently.
// This is O(n * (n + m)) which may be too slow for large n.
// For the IOI solution, the optimized merging is needed.
// Let me implement the straightforward approach for correctness:
fill(p.begin(), p.end(), 0);
for (int start = 0; start < n; start++) {
set<int> keys;
vector<bool> vis(n, false);
queue<int> q;
set<int> key_set;
auto add_room = [&](int room) {
if (vis[room]) return;
vis[room] = true;
q.push(room);
int key = r[room];
if (!key_set.count(key)) {
key_set.insert(key);
}
};
add_room(start);
bool changed = true;
while (changed) {
changed = false;
// Process BFS queue
while (!q.empty()) {
int room = q.front(); q.pop();
for (int j : adj[room]) {
if (key_set.count(c[j])) {
int other = u[j] == room ? v[j] : u[j];
if (!vis[other]) {
add_room(other);
changed = true;
}
}
}
}
// Check if new keys unlock new corridors
if (changed) continue;
for (int room_idx = 0; room_idx < n; room_idx++) {
if (!vis[room_idx]) continue;
for (int j : adj[room_idx]) {
if (key_set.count(c[j])) {
int other = u[j] == room_idx ? v[j] : u[j];
if (!vis[other]) {
add_room(other);
changed = true;
}
}
}
}
}
int cnt = 0;
for (int i = 0; i < n; i++) if (vis[i]) cnt++;
p[start] = cnt;
}
int min_p = *min_element(p.begin(), p.end());
vector<int> ans(n, 0);
for (int i = 0; i < n; i++)
if (p[i] == min_p) ans[i] = 1;
return ans;
}
int main() {
int n, m;
scanf("%d %d", &n, &m);
vector<int> r(n), u(m), v(m), c(m);
for (int i = 0; i < n; i++) scanf("%d", &r[i]);
for (int j = 0; j < m; j++) scanf("%d %d %d", &u[j], &v[j], &c[j]);
auto ans = find_reachable(r, u, v, c);
for (int i = 0; i < n; i++)
printf("%d%c", ans[i], " \n"[i == n - 1]);
return 0;
}
Complexity Analysis
Naive approach: $O(n(n + m))$ per room, total $O(n^2(n + m))$. Too slow for large inputs.
Optimized approach (not shown above): Uses Union-Find with key-based component merging. The idea is that if room $i$ reaches room $j$ and vice versa, they have the same reachable set. Process using a global BFS with component merging: $O((n + m) \alpha(n))$ amortized, where $\alpha$ is the inverse Ackermann function.
Space: $O(n + m)$.
Note: The full-score solution uses an approach where we start BFS from all rooms simultaneously, merging components when mutual reachability is detected. This avoids redundant exploration and achieves near-linear time complexity.
Code
C++ solution used for this page.
#include <bits/stdc++.h>
using namespace std;
// IOI 2021 - Keys
// n rooms with keys, m corridors requiring specific keys. For each room i,
// compute p[i] = number of reachable rooms. Return which rooms achieve min p[i].
//
// Approach: BFS from each room, collecting keys and exploring corridors
// unlocked by collected keys. Repeat until no new rooms are discovered.
// (This is the correct O(n*(n+m)) approach; the full-score solution uses
// component merging for near-linear time.)
vector<int> find_reachable(vector<int> r, vector<int> u,
vector<int> v, vector<int> c) {
int n = (int)r.size(), m = (int)u.size();
// Adjacency list: corridors incident to each room
vector<vector<int>> adj(n);
for (int j = 0; j < m; j++) {
adj[u[j]].push_back(j);
adj[v[j]].push_back(j);
}
vector<int> p(n, 0);
for (int start = 0; start < n; start++) {
vector<bool> vis(n, false);
set<int> key_set;
queue<int> q;
auto add_room = [&](int room) {
if (vis[room]) return;
vis[room] = true;
q.push(room);
key_set.insert(r[room]);
};
add_room(start);
bool changed = true;
while (changed) {
changed = false;
// BFS: explore corridors from visited rooms using collected keys
while (!q.empty()) {
int room = q.front();
q.pop();
for (int j : adj[room]) {
if (key_set.count(c[j])) {
int other = (u[j] == room) ? v[j] : u[j];
if (!vis[other]) {
add_room(other);
changed = true;
}
}
}
}
// New keys may unlock previously blocked corridors
if (changed) continue;
// Re-scan all visited rooms for newly unlockable corridors
for (int room_idx = 0; room_idx < n; room_idx++) {
if (!vis[room_idx]) continue;
for (int j : adj[room_idx]) {
if (key_set.count(c[j])) {
int other = (u[j] == room_idx) ? v[j] : u[j];
if (!vis[other]) {
add_room(other);
changed = true;
}
}
}
}
}
int cnt = 0;
for (int i = 0; i < n; i++)
if (vis[i]) cnt++;
p[start] = cnt;
}
int min_p = *min_element(p.begin(), p.end());
vector<int> ans(n, 0);
for (int i = 0; i < n; i++)
if (p[i] == min_p) ans[i] = 1;
return ans;
}
int main() {
int n, m;
scanf("%d %d", &n, &m);
vector<int> r(n), u(m), v(m), c(m);
for (int i = 0; i < n; i++) scanf("%d", &r[i]);
for (int j = 0; j < m; j++) scanf("%d %d %d", &u[j], &v[j], &c[j]);
auto ans = find_reachable(r, u, v, c);
for (int i = 0; i < n; i++)
printf("%d%c", ans[i], " \n"[i == n - 1]);
return 0;
}
Source Files and Assets
Raw files are still available here when you want the original TeX, C++, or statement assets.