Mushrooms
There are n mushrooms, each either type A or type B. Mushroom 0 is known to be type A. You can query a function use\_machine(x) that takes a sequence of mushroom indices and returns the number of adjacent pairs in the...
Problem Statement
Rendered from the "Problem Summary" section in the LaTeX write-up.
There are $n$ mushrooms, each either type A or type B. Mushroom 0 is known to be type A. You can query a function use_machine(x) that takes a sequence of mushroom indices and returns the number of adjacent pairs in the sequence that are of different types. Determine the total count of type A mushrooms using at most a limited number of queries (around $O(\sqrt{n})$ or $O(n^{2/3})$).
Editorial
The solution write-up is rendered from the LaTeX source, with equations kept live through MathJax.
Solution Approach
Basic Observations
use_machine({0, i})returns 0 if $i$ is type A, 1 if type B. This identifies one mushroom per query.To be more efficient, we can batch multiple unknowns into a single query.
Batch Queries
If we know $k$ mushrooms of type A (say $a_1, a_2, \ldots, a_k$) and want to classify $k$ unknowns $u_1, \ldots, u_k$, we interleave them:
\texttt{use_machine({$a_1, u_1, a_2, u_2, \ldots, a_k, u_k$})}
The result tells us the total number of type-changes. Since the known mushrooms are all type A, each transition $a_i \to u_i$ contributes 1 if $u_i$ is B, and each transition $u_i \to a_{i+1}$ contributes 1 if $u_i$ is B. So each unknown $u_i$ (except possibly the last) contributes 0 or 2 to the count, and we can deduce the types.
More precisely, arranging as $[a_1, u_1, a_2, u_2, \ldots]$:
Result $r$. Since A-A transition contributes 0 and A-B or B-A contributes 1:
Each B-type unknown contributes exactly 2 to the count (one transition going in, one going out), except the last unknown which contributes 1.
This allows classifying $k-1$ unknowns per query, given $k$ known type-A mushrooms.
Strategy
Start by identifying the first few mushrooms individually.
Once we have a pool of known type-A (or type-B) mushrooms of size $k$, batch-classify $k$ unknowns per query.
As our pool grows, each query becomes more efficient.
Total queries: approximately $O(\sqrt{n})$.
C++ Solution
#include <bits/stdc++.h>
using namespace std;
int use_machine(vector<int> x); // provided by grader
int count_mushrooms(int n) {
if (n == 1) return 1;
if (n == 2) {
return 1 + (1 - use_machine({0, 1}));
}
vector<int> A = {0}; // known type A
vector<int> B; // known type B
// Identify mushrooms 1 and 2
int r1 = use_machine({0, 1});
int r2 = use_machine({0, 2});
if (r1 == 0) A.push_back(1); else B.push_back(1);
if (r2 == 0) A.push_back(2); else B.push_back(2);
int countA = A.size();
int idx = 3;
while (idx < n) {
// Use whichever pool is larger
vector<int>& pool = (A.size() >= B.size()) ? A : B;
bool using_A = (&pool == &A);
int k = pool.size();
int batch = min(k, n - idx); // number of unknowns to classify
// Build interleaved sequence
vector<int> seq;
for (int i = 0; i < batch; i++) {
seq.push_back(pool[i]);
seq.push_back(idx + i);
}
// Optionally add one more known at the end for the last unknown
if (batch <= k) {
// The sequence is [p0, u0, p1, u1, ..., p_{batch-1}, u_{batch-1}]
// Transitions: p0-u0, u0-p1, p1-u1, ..., p_{batch-1}-u_{batch-1}
// Total transitions = 2 * (number of B-type unknowns if using_A pool)
// ... well, let's compute carefully.
}
int r = use_machine(seq);
// Decode: seq = [p0, u0, p1, u1, ..., p_{b-1}, u_{b-1}]
// Transitions between consecutive elements:
// (p0,u0), (u0,p1), (p1,u1), (u1,p2), ...
// If using A pool: all p_i are type A.
// Transition (p_i, u_i): 1 if u_i is B, 0 if A.
// Transition (u_i, p_{i+1}): 1 if u_i is B, 0 if A.
// So each u_i contributes 2 if B (or 0 if A), except u_{batch-1}
// which only has the left transition (p_{batch-1}, u_{batch-1}).
// Total r = 2 * (# B among u_0..u_{batch-2}) + (u_{batch-1} is B ? 1 : 0)
// Determine type of last unknown
int last_contrib = r % 2; // 1 if last unknown is B (when using A pool)
int rest_B = (r - last_contrib) / 2;
// But we don't know which specific unknowns (besides last) are B.
// We only know the total count.
// Actually, we CAN determine each unknown's type individually if we
// use a smarter interleaving. Let me reconsider.
// Alternative: use [p0, u0, p1, u1, ...] and process prefix queries.
// But that uses too many queries.
// For the batch approach, we ONLY learn the total count, not individual types.
// So we add all unknowns to a "pending" list and only count.
// Correction: we do learn each type! Consider:
// seq = [p0, u0, p1, u1, p2, u2, ...]
// r = sum of |type(seq[i]) - type(seq[i+1])| for adjacent pairs
// The transitions are:
// t0 = |A - type(u0)|
// t1 = |type(u0) - A|
// t2 = |A - type(u1)|
// t3 = |type(u1) - A|
// ...
// So t_{2i} + t_{2i+1} = 2 * (u_i is B), for i < batch-1
// And t_{2*(batch-1)} = (u_{batch-1} is B)
// But we only get the sum r, not individual t values!
// So we know: rest_B = total B unknowns among u_0..u_{batch-2}
// and whether u_{batch-1} is B.
// We DON'T know which specific ones are B among u_0..u_{batch-2}.
// For counting purposes, this is sufficient!
// We just need the total count of A mushrooms, not their identities.
// Count A types in this batch:
int batchA, batchB;
if (using_A) {
// last unknown: B if last_contrib=1, A if 0
batchB = rest_B + last_contrib;
batchA = batch - batchB;
} else {
// Using B pool: transitions flip
// t_{2i} + t_{2i+1} = 2 * (u_i is A), etc.
int rest_A = (r - last_contrib) / 2;
int lastA = last_contrib;
batchA = rest_A + lastA;
batchB = batch - batchA;
}
countA += batchA;
// We need to grow our pools. We can't add unknowns since we don't know
// their individual types. So we sacrifice query efficiency for pool growth.
// Better approach: for the first few, identify individually to grow pools.
// Then batch for counting.
// Actually, let's refine: process unknowns ONE AT A TIME when pool is small,
// and batch when pool is large enough.
// For simplicity, let me just use individual queries for small pool
// and batch for large.
idx += batch;
}
// Actually, the above logic mixes batching strategies. Let me rewrite cleanly.
// (The code below is a clean reimplementation.)
return countA;
}
// Clean reimplementation:
int count_mushrooms_v2(int n) {
if (n <= 2) {
if (n == 1) return 1;
return 1 + (1 - use_machine({0, 1}));
}
vector<int> known[2]; // known[0] = type A indices, known[1] = type B indices
known[0].push_back(0);
int countA = 1;
int idx = 1;
while (idx < n) {
// Use the larger pool
int t = (known[1].size() > known[0].size()) ? 1 : 0;
int k = known[t].size();
if (k < 2 || n - idx < 2) {
// Identify one at a time
int r = use_machine({known[t][0], idx});
int type_idx = (r == 0) ? t : (1 - t);
if (type_idx == 0) countA++;
known[type_idx].push_back(idx);
idx++;
} else {
// Batch: interleave known[t] with unknowns
int batch = min(k, n - idx);
vector<int> seq;
for (int i = 0; i < batch; i++) {
seq.push_back(known[t][i % k]);
seq.push_back(idx + i);
}
int r = use_machine(seq);
// Last unknown: contributes r%2 transitions
int last_diff = r % 2;
int rest_diff = (r - last_diff) / 2;
// rest_diff = number of unknowns (among first batch-1) that differ from type t
// last_diff = 1 if last unknown differs from type t
int diff_count = rest_diff + last_diff;
int same_count = batch - diff_count;
if (t == 0) {
countA += same_count;
} else {
countA += diff_count;
}
// We know the type of the last unknown; add it to pool
int last_type = last_diff ? (1 - t) : t;
known[last_type].push_back(idx + batch - 1);
// For batch-1 unknowns, we know total but not individual types
// We can't add them to pools. That's fine for counting.
idx += batch;
}
}
return countA;
}
int main() {
int n;
scanf("%d", &n);
printf("%d\n", count_mushrooms_v2(n));
return 0;
}
Complexity Analysis
Query complexity: Initially, we use $O(1)$ query per mushroom. Once the pool reaches size $k$, each query classifies $k$ mushrooms. The pool grows by 1 per batch, so after $k$ batches, pool has size $k + k$ and we classify $k^2$ mushrooms total. Overall queries $\approx O(\sqrt{n})$.
Time complexity: $O(n)$ for processing all mushrooms.
Space complexity: $O(n)$.
Code
C++ solution used for this page.
#include <bits/stdc++.h>
using namespace std;
// IOI 2020 - Mushrooms
// Determine the count of type-A mushrooms using limited machine queries.
// Strategy: grow known pools, then batch-classify using interleaving.
// Mushroom 0 is known type A. use_machine returns number of adjacent
// different-type pairs in the sequence.
int use_machine(vector<int> x); // provided by grader
int count_mushrooms(int n) {
if (n == 1) return 1;
if (n == 2) return 1 + (1 - use_machine({0, 1}));
// known[0] = type A indices, known[1] = type B indices
vector<int> known[2];
known[0].push_back(0);
int countA = 1;
int idx = 1;
while (idx < n) {
// Use the larger pool for batching
int t = ((int)known[1].size() > (int)known[0].size()) ? 1 : 0;
int k = (int)known[t].size();
if (k < 2 || n - idx < 2) {
// Identify one mushroom at a time
int r = use_machine({known[t][0], idx});
int type_idx = (r == 0) ? t : (1 - t);
if (type_idx == 0) countA++;
known[type_idx].push_back(idx);
idx++;
} else {
// Batch: interleave known[t] mushrooms with unknowns
int batch = min(k, n - idx);
vector<int> seq;
for (int i = 0; i < batch; i++) {
seq.push_back(known[t][i]);
seq.push_back(idx + i);
}
int r = use_machine(seq);
// Decode transitions:
// Each unknown (except last) is flanked by two known mushrooms,
// contributing 2 to r if different from type t, 0 if same.
// Last unknown contributes 1 if different, 0 if same.
int last_diff = r % 2;
int rest_diff = (r - last_diff) / 2;
int diff_count = rest_diff + last_diff;
int same_count = batch - diff_count;
if (t == 0) {
countA += same_count; // same as type A
} else {
countA += diff_count; // different from type B = type A
}
// We know the last unknown's type; add it to the pool
int last_type = last_diff ? (1 - t) : t;
known[last_type].push_back(idx + batch - 1);
idx += batch;
}
}
return countA;
}
int main() {
int n;
scanf("%d", &n);
printf("%d\n", count_mushrooms(n));
return 0;
}
Source Files and Assets
Raw files are still available here when you want the original TeX, C++, or statement assets.