题目

给定一个整数 n 和一个 无重复 黑名单整数数组 blacklist 。设计一种算法，从 [0, n - 1] 范围内的任意整数中选取一个 未加入 黑名单 blacklist 的整数。任何在上述范围内且不在黑名单 blacklist 中的整数都应该有 同等的可能性 被返回。

优化你的算法，使它最小化调用语言 内置 随机函数的次数。

实现 Solution 类:

• Solution(int n, int[] blacklist) 初始化整数 n 和被加入黑名单 blacklist 的整数
• int pick() 返回一个范围为 [0, n - 1] 且不在黑名单 blacklist 中的随机整数

示例

示例 1：

输入 [“Solution”, “pick”, “pick”, “pick”, “pick”, “pick”, “pick”,

“pick”] [[7, [2, 3, 5]], [], [], [], [], [], [], []]

输出 [null, 0, 4,1, 6, 1, 0, 4]

解释 Solution solution = new Solution(7, [2, 3, 5]); solution.pick(); // 返回0，任何[0,1,4,6]的整数

都可以。注意，对于每一个pick的调用，// 0、1、4和6的返回概率必须相等(即概率为1/4).

solution.pick(); // 返回 4 solution.pick(); // 返回 1

solution.pick(); // 返回 6 solution.pick(); // 返回 1 solution.pick(); // 返回 0 solution.pick(); // 返回 4

提示:

1 <= n <= 109

0 <= blacklist.length <= min(105, n - 1)

0 <= blacklist[i] < n

blacklist 中所有值都 不同

pick 最多被调用 2 * 104 次

思路

借鉴了评论区大佬的思路

思路：

• 黑名单长度为s，我们从[0, N-s)中取随机值，这个随机值有可能在黑名单中，怎么办？
• [0, N-s)内的元素，如果有i个在黑名单中，那么在[N-s, N)中，必定有i个元素不在黑名单中
• 对[0, N-s)中的黑名单元素和[N-s, N)中不在黑名单中的元素做映射m，必定可以一一对应，怎么对应倒是无所谓
• 从[0, N-s)中取随机值r，如果r不在黑名单中，直接返回；如果r在黑名单中，则m[r]一定不在黑名单，返回m[r]

题解

from random import randint
class Solution:
    def __init__(self, N, blacklist):
        """
        :type N: int
        :type blacklist: List[int]
        """
        self.s = N - len(blacklist)
        # 小于s的黑名单元素集合
        b_lt_s = {i for i in blacklist if i < self.s}
        # 大于s的非黑名单元素集合
        # 等价于：w_gt_s = {i for i in range(self.s, N)} - set(blacklist)，感谢评论
        w_gt_s = {*range(self.s, N)} - set(blacklist)
        # 做映射，使用zip方便一点
        # 等价于：self.m = {k: v for k,v in zip(b_lt_s, w_gt_s)}
        self.m = dict(zip(b_lt_s, w_gt_s))
    def pick(self):
        """
        :rtype: int
        """
        r = randint(0, self.s-1)
        return r if r not in self.m else self.m[r]