def__init__(self, n: int, blacklist: List[int]): self.r = n m = len(blacklist) if m <= n / 2: self.type = 0# 拒绝采样 self.b = set(blacklist) else: self.type = 1 self.b = set(range(n)) - set(blacklist)
defpick(self) -> int: ifself.type == 0: whileTrue: val = random.randint(0, self.r - 1) if val notinself.b: return val else: val = self.b.pop() self.b.add(val) return val
// Code language: C++ classSolution { public: bool type; unordered_set<int> b; int n; vector<int> whitelist;
Solution(int n, vector<int>& blacklist) { this->n = n; int m = blacklist.size(); if (m <= n / 2) type = true; // 拒绝采样 else type = false; // 枚举
for (int val : blacklist) b.emplace(val); if (!type) { for (int i = 0; i < n; ++i) { if (b.find(i) == b.end()) whitelist.emplace_back(i); } } } intpick(){ if (type) { while (true) { int val = rand() % n; if (b.find(val) == b.end()) return val; } } int m = whitelist.size(); int val = rand() % m; return whitelist[val]; } };
时间复杂度: 预处理O(m),
其中拒绝采样的预处理是 O(m), 枚举的预处理是 O(n) 但枚举的情况 n, m
是一个数量级,所以也可以认为是 O(m),pick函数为 O(1),
其中拒绝采样期望是 O(2),
枚举是 O(1)
空间复杂度: O(m)
哈希映射
参考官解的做法,实际上可取的值为 n − m 个,将黑名单的 m 个数值映射为 n − m ~ n 中的数字即可。
通俗的说,我们将 [0, n − m)
中在黑名单的数字拿掉,然后用 [n − m, n)
中不在黑名单中的数字填上,但我们不能列出完整的 [0, n − m)
数组,所以只用哈希表记录那些替换掉的即可。
def__init__(self, n: int, blacklist: List[int]): self.m = n - len(blacklist) self.st = dict() s = set(blacklist) x = self.m for b in blacklist: if b < self.m: while x in s: x += 1 self.st[b] = x x += 1
defpick(self) -> int: val = random.randint(0, self.m - 1) if val inself.st: returnself.st[val] return val