leetcode-go/leetcode/hard/710.cc

#include <algorithm>
#include <cstdlib>
#include <ctime>
#include <set>
#include <vector>

#include <iostream>

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<int>& ban;
    set<int> 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<int>& blacklist) : N(N), MAX(N - blacklist.size()), ban(blacklist), ban_set(set<int>(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<int> 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;
}