710. 黑名单中的随机数

exiaohu 于 2022-06-26 发布

题目链接: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)