First, let's determine who is originally invited and what criteria they are invited under. This takes $O(C+\sum n_i)$ time if we loop over each criterion in increasing order, and then the contestants satisfying that criterion in increasing order.
Then we can process the declinations one at a time. When an invited contestant declines, that frees up one slot for the criterion $c$ they were invited under, which results in one of the following cases:
To process all of these cases efficiently, let's maintain a queue for each criterion with the contestants who satisfy that criterion but haven't been invited yet, in increasing order. Then, when processing criterion $c$, we can pop contestants from the front of the corresponding queue who have already declined or are currently invited under some criterion $c'<c$ until we reach one of the three cases.
The runtime when processing a declination is bounded by a constant plus the number of queue elements popped. Therefore, the total runtime is bounded by $C$ plus the initial sum of queue sizes, which is $O(C+\sum n_i)$.
#include <bits/stdc++.h>
using namespace std;
template <class T> using V = vector<T>;
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
int N, C;
cin >> N >> C;
V<int> f(C);
for (int &t : f) cin >> t;
V<int> p(N);
for (int &t : p) {
cin >> t;
--t; // zero-index
}
V<queue<int>> q(C);
V<int> invited_under(N, C);
// zero-indexed criterion that each contestant is invited under
// also, C = not invited, -1 = declined
for (int i = 0; i < N; ++i) {
int n;
cin >> n;
for (int j = 0; j < n; ++j) {
int c;
cin >> c;
--c;
q.at(c).push(i);
}
}
// initial invitations
int64_t ans = 0;
for (int c = 0; c < C; ++c) {
int invited_under_c = 0;
while (size(q.at(c)) && invited_under_c < f.at(c)) {
int x = q.at(c).front();
q.at(c).pop();
if (invited_under.at(x) < c) continue;
invited_under.at(x) = c;
ans += x + 1;
++invited_under_c;
}
}
// process declinations one at a time
for (int d : p) {
cout << ans << "\n";
int c = invited_under.at(d);
invited_under.at(d) = -1;
if (c == C) continue;
ans -= d + 1;
while (size(q.at(c))) {
int x = q.at(c).front();
q.at(c).pop();
if (invited_under.at(x) < c) continue;
assert(invited_under.at(x) > c);
swap(c, invited_under.at(x));
if (c == C) {
ans += x + 1;
break;
}
}
}
}