IOI 2021
IOI 2021

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...

Updated May 21, 2026
Track IOI
Year 2021
Statement Rendered from TeX
TeXC++Rendered statement

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:

  1. For each room $i$, start a BFS/DFS exploration.

  2. Maintain the set of collected keys and frontier corridors.

  3. When a new key is collected, check all corridors requiring that key type.

  4. 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:

  1. Start BFS from room 0. Maintain key set $K$, visited rooms $V$, and pending corridors $E$.

  2. When a new key of type $t$ is added to $K$, add all corridors of type $t$ incident to visited rooms to the frontier.

  3. If BFS terminates (no more reachable rooms), $p[0] = |V|$.

  4. 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.

C++

Clean code view with a raw-file link when you want the original source.

Raw file
#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.

Show raw files