The input is two lists: the subroutines called and their arguments. Solution's constructor has two arguments, N and the blacklist B. pick has no arguments. Arguments are always wrapped with a list, even if there aren't any.
Solution
// Solution 1: Binary Search. O(log)
class Solution {
vector<int> v;
std::random_device rd;
int mod;
public:
Solution(int N, vector<int> blacklist) {
mod = N - blacklist.size();
sort(blacklist.begin(), blacklist.end());
v = blacklist;
v.push_back(N);
for (int i = 0; i < v.size(); i++) v[i] -= i;
}
int pick() {
int index = rd() % mod;
auto it = upper_bound(v.begin(), v.end(), index);
return index + (it - v.begin());
}
};
// Solution 2: hash_map; O(1)
class Solution {
random_device rd;
unordered_map<int, int> updated_index;
int mod;
public:
Solution(int N, vector<int>& blacklist) {
mod = N - blacklist.size();
set<int> blacklist_set(blacklist.begin(), blacklist.end());
int swap_num = N - 1;
for (int blacklist_element : blacklist_set) {
if (blacklist_element >= mod) break;
while (blacklist_set.count(swap_num) > 0) {
swap_num--;
}
updated_index[blacklist_element] = swap_num;
swap_num--;
}
}
int pick() {
int index = rd() % mod;
if (updated_index.count(index) > 0) {
return updated_index[index];
}
return index;
}
};