IOI 2021 - Keys

#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;
}

\documentclass[11pt,a4paper]{article}
\usepackage[utf8]{inputenc}
\usepackage[T1]{fontenc}
\usepackage{amsmath,amssymb,amsthm}
\usepackage{listings}
\usepackage{geometry}
\usepackage{hyperref}
\geometry{margin=2.5cm}

\lstset{
  language=C++,
  basicstyle=\ttfamily\small,
  keywordstyle=\bfseries,
  breaklines=true,
  numbers=left,
  numberstyle=\tiny,
  frame=single,
  tabsize=2
}

\title{IOI 2021 -- Keys}
\author{Solution}
\date{}

\begin{document}
\maketitle

\section{Problem Summary}
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]$.

\section{Solution Approach}

\subsection{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))$.

\subsection{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:

\begin{enumerate}
  \item For each room $i$, start a BFS/DFS exploration.
  \item Maintain the set of collected keys and frontier corridors.
  \item When a new key is collected, check all corridors requiring that key type.
  \item If exploration from room $i$ reaches room $j$, and we previously explored from $j$ and found its reachable set, merge them.
\end{enumerate}

\subsection{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:
\begin{enumerate}
  \item Start BFS from room 0. Maintain key set $K$, visited rooms $V$, and pending corridors $E$.
  \item When a new key of type $t$ is added to $K$, add all corridors of type $t$ incident to visited rooms to the frontier.
  \item If BFS terminates (no more reachable rooms), $p[0] = |V|$.
  \item If during BFS we enter a room already in another component, merge components.
\end{enumerate}

\section{C++ Solution}

\begin{lstlisting}
#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;
}
\end{lstlisting}

\section{Complexity Analysis}
\begin{itemize}
  \item \textbf{Naive approach:} $O(n(n + m))$ per room, total $O(n^2(n + m))$. Too slow for large inputs.
  \item \textbf{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.
  \item \textbf{Space:} $O(n + m)$.
\end{itemize}

\textbf{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.

\end{document}