题目链接:710. 黑名单中的随机数
维护一个白名单列表的范围。
取样时先以范围大小为权重取一个范围,然后在这个范围里均匀采样。
采样次数为 2。
class Solution:
def __init__(self, n: int, blacklist: List[int]):
assert all(0 <= v <= n - 1 for v in blacklist)
blacklist = sorted(blacklist)
whitelist_range = []
if blacklist:
if blacklist[0] > 0:
whitelist_range.append((0, blacklist[0] - 1))
if blacklist[-1] < n - 1:
whitelist_range.append((blacklist[-1] + 1, n - 1))
for l, r in zip(blacklist[:-1], blacklist[1:]):
if l < r - 1:
whitelist_range.append((l + 1, r - 1))
else:
whitelist_range.append((0, n - 1))
self.whitelist_range = whitelist_range
self.weights = [r - l + 1 for l, r in whitelist_range]
def pick(self) -> int:
l, r = random.choices(population=self.whitelist_range, weights=self.weights, k=1)[0]
return random.randint(l, r)