LeetCode710. 黑名单中的随机数

题目描述

给定一个整数 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 * 10^4 次

题解1 - 前缀和 + 二分

此解法在思路上类似某一天的每日一题『497. 非重叠矩形中的随机点』，该题题解。
将范围 $[0,n-1]$ 根据黑名单中的数分为若干段，假设白名单里的数（不在黑名单里的数）个数为 $len$，在 $[1,len]$ 中随机取值 $v$，表示白名单中第 $v$ 个数，查找这个数在哪段中，即可根据该段的端点求出这个值是多少。
分段操作要求黑名单数组有序，统计每段中含有多少个数字，使用前缀和表示，前缀和天然带有单调性，可以使用二分查找第 $v$ 个数。

代码

class Solution {
private:
    mt19937 gen{ random_device{}() };
    vector<pair<int, int>> regions;
    vector<int> prefixSum;
    int len;

public:
    Solution(int n, vector<int>& blacklist) {
        int m = blacklist.size();
        len = n - m;
        sort(blacklist.begin(), blacklist.end());
        if (m == 0) {
            regions.emplace_back(0, n - 1);
        } else {
            if (blacklist[0] != 0) {
                regions.emplace_back(0, blacklist[0] - 1);
            }
            for (int i = 1; i < m; ++i) {
                regions.emplace_back(blacklist[i - 1] + 1, blacklist[i] - 1);
            }
            if (blacklist[m - 1] != n - 1) {
                regions.emplace_back(blacklist[m - 1] + 1, n - 1);
            }
        }
        prefixSum.emplace_back(0);
        for (const auto& region: regions) {
            prefixSum.emplace_back(prefixSum.back() + region.second - region.first + 1);
        }
    }
    
    int pick() {
        uniform_int_distribution<int> dis(1, len);
        int v = dis(gen);
        int idx = lower_bound(prefixSum.begin(), prefixSum.end(), v) - prefixSum.begin();
        auto p = regions[idx - 1];
        return p.second - (prefixSum[idx] - v);
    }
};

复杂度分析

时间复杂度：构造函数的复杂度为 $O(m\log{m})$，pick() 操作的复杂度 $O(\log{m})$
空间复杂度：$O(m)$

题解2 - 哈希映射

设黑名单长度为 $m$，想要等概率选取，最理想的情况为白名单里的数正好在 $[0,n-m-1]$ 中，考虑出现在 $[0,n-m-1]$ 中的黑名单数字，将它们映射到 $[n-m,n-1]$ 中的白名单数字，即可直接在 $[0,n-m-1]$ 随机取数。
为了方便查找，将 $[n-m,n-1]$ 中的黑名单数字放到一个哈希集合中。

代码

class Solution {
private:
    mt19937 gen{ random_device{}() };
    unordered_map<int, int> mp;
    int len;

public:
    Solution(int n, vector<int>& blacklist) {
        int m = blacklist.size();
        len = n - m;
        unordered_set<int> blackset;
        for (const auto& b: blacklist) {
            if (b >= len) {
                blackset.emplace(b);
            }
        }
        int i = len;
        for (const auto& b: blacklist) {
            if (b < len) {
                while (blackset.find(i) != blackset.end()) {
                    ++i;
                }
                mp[b] = i++;
            }
        }
    }
    
    int pick() {
        uniform_int_distribution<int> dis(0, len - 1);
        int idx = dis(gen);
        return mp.find(idx) == mp.end() ? idx : mp[idx];
    }
};

可以将黑名单数组排序通过二分法求出大于等于 $n-m$ 的最小值，省去判断，但复杂度会增加。

class Solution {
private:
    mt19937 gen{ random_device{}() };
    unordered_map<int, int> mp;
    int len;

public:
    Solution(int n, vector<int>& blacklist) {
        int m = blacklist.size();
        len = n - m;

        sort(blacklist.begin(), blacklist.end());
        int bound = lower_bound(blacklist.begin(), blacklist.end(), len) - blacklist.begin();
        unordered_set<int> blackset;
        
        for (int i = bound; i < m; ++i) {
            blackset.emplace(blacklist[i]);
        }
        int t = len;
        for (int i = 0; i < bound; ++i) {
            while (blackset.find(t) != blackset.end()) {
                ++t;
            }
            mp[blacklist[i]] = t++;
        }
    }
    
    int pick() {
        uniform_int_distribution<int> dis(0, len - 1);
        int idx = dis(gen);
        return mp.find(idx) == mp.end() ? idx : mp[idx];
    }
};