【哈希表, 黑名单映射】黑名单中的随机数

题目

710. 黑名单中的随机数

---

给定一个整数 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 <= 10^9
• 0 <= blacklist.length <= min(10^5, n - 1)
• 0 <= blacklist[i] < n
• blacklist 中所有值都 不同
• pick 最多被调用 2 * 10^4 次

解题

方法一：黑名单映射

思路

设 blacklist 的长度为 m。

考察一个特殊的例子：所有黑名单数全部在区间 $[n-m,n)$ 范围内。此时我们可以直接在 $[0,n-m)$ 范围内取随机整数。

这给我们一个启示，对于在 $[0,n-m)$ 范围内的黑名单数，我们可以将其映射到 $[n-m,n)$ 范围内的非黑名单数（白名单数）上。每次 pick() 时，仍然可以在 $[0,n−m)$ 范围内取随机整数（设其为 x），那么：

• 如果 x 不在黑名单中，则直接返回 x；
• 如果 x 在黑名单中，则返回 x 映射到 $[n-m,n)$ 范围内的白名单数。

我们可以在初始化时，构建一个从 $[0,n-m)$ 范围内的黑名单数到 $[n-m,n)$ 的白名单数的映射：

1. 将 $[n-m,n)$ 范围内的黑名单数存入一个哈希集合 black；
2. 初始化白名单数 $w=n−m$；
3. 对于每个 $[0,n-m)$ 范围内的黑名单数 b，首先不断增加 w 直至其不在黑名单中，然后将 b 映射到 w 上，并将 w 自增。

代码

class Solution {
    private Random rand;
    private int bound;
    private Map<Integer, Integer> map;

    public Solution(int n, int[] blacklist) {
        rand = new Random();
        bound = n - blacklist.length;
        Set<Integer> blacks = new HashSet<Integer>() {{
            for (int b : blacklist) {
                if (b >= bound) this.add(b);
            }
        }};
        map = new HashMap<>();
        int num = bound;
        for (int b : blacklist) {
            if (b < bound) {
                while (blacks.contains(num)) num++;
                map.put(b, num);
                num++;
            }
        }
    }
    
    public int pick() {
        int randNum = rand.nextInt(bound);
        return map.getOrDefault(randNum, randNum);
    }
}

class Solution {
    int bound;
    unordered_map<int, int> mp;

public:
    Solution(int n, vector<int>& blacklist) {
        bound = n - blacklist.size();
        unordered_set<int> rear_blacks;
        for (int& black : blacklist) {
            if (black >= bound) rear_blacks.insert(black);
        }
        int rear_white = bound;
        for (int& black : blacklist) {
            if (black < bound) {
                while (rear_blacks.find(rear_white) != rear_blacks.end()) ++rear_white;
                mp[black] = rear_white++;
            }
        }
    }
    
    int pick() {
        int rand_num = rand() % bound;
        return mp.find(rand_num) == mp.end() ? rand_num : mp[rand_num];
    }
};