Shared Mushroom Baskets
National Olympiad — National round (2025 season)
Submission name: mushrooms.cpp
Problem summary
M basket lists record which mushrooms have been in each basket. For each pair of distinct mushrooms, find the sum of the one-based indices of baskets containing both.
Input
N M; for baskets 1 through M, their size followed by their mushroom labels; Q; then Q distinct-label pairs A B.
Output
One sum per query, in input order.
Constraints
1 ≤ N,S,Q ≤ 10⁵, where S is the total basket-list length; 2 ≤ basket size ≤ N; 1 ≤ M ≤ S/2.
Worked example
An independently prepared example for this explanation.
Input
6 3 3 1 2 3 2 5 6 3 2 5 6 4 1 2 2 5 3 6 5 6
Output
1 3 0 5
Solution analysis
- Build the sorted list of basket indices containing each mushroom. Classify a mushroom as heavy when it appears in more than B≈sqrt(S) baskets. There are at most S/B heavy mushrooms. Two light lists can be intersected by a linear two-pointer merge in O(B) time, adding the common basket indices.
- If a query contains a heavy mushroom, assign the query to one such endpoint. Process these heavy endpoints one at a time. For heavy h, clear an accumulator indexed by mushroom, visit each basket containing h, and add that basket’s index to every mushroom in it. The accumulator now answers all queries assigned to h in constant time each.
- Only one accumulator is stored, avoiding a potentially large table of all heavy-to-mushroom answers. In the example mushrooms 5 and 6 share baskets 2 and 3, so their answer is 5. A pair with no common basket gets zero naturally.
Why it works
Basket lists are constructed in increasing basket order, so the two-pointer intersection returns exactly all common indices for light pairs. For a heavy endpoint h, each basket containing h contributes its index once to exactly the mushrooms that share it, so the accumulator equals the requested common-index sum. Every query is handled by one of these exact methods, preserving its original output position.
Watch out. Heavy means frequent across baskets, not a large basket. Attach a query with two heavy endpoints to only one. Use 64-bit sums of basket indices.
Commented C++17 solution
An independently written complete program with standard headers.
#include<algorithm>
#include<cmath>
#include<iostream>
#include<utility>
#include<vector>
// Algorithm: Build the sorted list of basket indices containing each mushroom.
using namespace std;
int main(){
ios::sync_with_stdio(false);
cin.tie(nullptr);
int n,m;
cin>>n>>m;
vector<vector<int>>baskets(m+1),memberships(n+1);
int total=0;
for(int basket=1;basket<=m;++basket){
int size;
cin>>size;
total+=size;
baskets[basket].resize(size);
for(int&mushroom:baskets[basket]){
cin>>mushroom;
memberships[mushroom].push_back(basket);
}
}
int threshold=max(1,static_cast<int>(sqrt(total)));
int q;
cin>>q;
vector<long long>answer(q),accumulator(n+1);
vector<vector<pair<int,int>>>assigned(n+1);
for(int query=0;query<q;++query){
int a,b;
cin>>a>>b;
if(static_cast<int>(memberships[a].size())>threshold){
assigned[a].push_back({b,query});
}
else if(static_cast<int>(memberships[b].size())>threshold){
assigned[b].push_back({a,query});
}
else{
size_t i=0;
size_t j=0;
while(i<memberships[a].size()&&j<memberships[b].size()){
if(memberships[a][i]==memberships[b][j]){
answer[query]+=memberships[a][i];
++i;
++j;
}
else if(memberships[a][i]<memberships[b][j]){
++i;
}
else{
++j;
}
}
}
}
for(int heavy=1;heavy<=n;++heavy){
if(assigned[heavy].empty()){
continue;
}
fill(accumulator.begin(),accumulator.end(),0);
for(int basket:memberships[heavy]){
for(int mushroom:baskets[basket]){
accumulator[mushroom]+=basket;
}
}
for(auto query:assigned[heavy]){
answer[query.second]=accumulator[query.first];
}
}
for(long long value:answer){
cout<<value<<'\n';
}
}