#include #include #include #include #include #include using std::set; using std::vector; using std::sort; using std::cout; using std::endl; class Solution { private: const int N; const int MAX; vector& ban; set ban_set; int random() { return rand() % MAX; } int search_offset(int n) { int beg = -1, end = ban.size() - 1; while (end - beg > 1) { int mid = beg + (end - beg) / 2; if (ban.at(mid) < n) { beg = mid; } else { end = mid; } } return n < ban.at(end) ? end : end + 1; } public: Solution(int N, vector& blacklist) : N(N), MAX(N - blacklist.size()), ban(blacklist), ban_set(set(blacklist.begin(), blacklist.end())) { srand(time(0)); sort(ban.begin(), ban.end()); } int pick() { int r = random(); if (N == MAX) { return r; } // Binary search candidates from 0 to N - 1 int beg = 0, end = N - 1; while(beg <= end) { int mid = beg + (end - beg) / 2; // To judge candidate is available or not: // 1. Not banned; 2. Offset is correct int rr = mid - search_offset(mid); if (rr < r) { beg = mid + 1; } else if (r < rr) { end = mid - 1; } else { if (ban_set.find(mid) == ban_set.end()) { return mid; } else { end = mid - 1; } } } return -1; } }; /** * Your Solution object will be instantiated and called as such: * Solution* obj = new Solution(N, blacklist); * int param_1 = obj->pick(); */ // 0 1 2 3 // 2 // 0 4 5 int main() { int N = 4; vector blacklist; blacklist.push_back(2); Solution* obj = new Solution(N, blacklist); int param_1 = obj->pick(); int param_2 = obj->pick(); int param_3 = obj->pick(); cout << param_1 << ", " << param_2 << ", " << param_3 << endl; return 0; }