题目地址:
https://leetcode.com/problems/random-pick-with-blacklist/description/
有一个长 n n n的数组 A A A,再给定一个数组 B B B,其表示 A A A中的若干个下标的黑名单。要求实现一个方法,该方法获取黑名单下标之外的随机的 0 ∼ n − 1 0\sim n-1 0∼n−1的下标,每个下标要等概率返回。
设 m = l B m=l_B m=lB,那么 A A A的前 n − m n-m n−m个数里如果有 x x x个下标在黑名单之中,那么A的后 m m m个数一定有 x x x个下标不在黑名单之中,我们可以在前 x x x个在黑名单之中的下标到后 x x x个不在黑名单的下标做一个映射。每次随机一个 [ 0 , n − m − 1 ] [0,n-m-1] [0,n−m−1]的数,如果不在黑名单则直接返回,否则映射一下再返回。代码如下:
class Solution {
public:
int n, len;
unordered_map<int, int> mp;
Solution(int n, vector<int>& blacklist) {
this->n = n;
len = blacklist.size();
unordered_set<int> st;
// 求出后x个不在黑名单中的数下标
for (int i = n - len; i < n; i++) st.insert(i);
for (int i : blacklist) st.erase(i);
// 做映射
auto it = st.begin();
for (int i : blacklist)
if (i < n - len) mp[i] = *it++;
}
int pick() {
int k = rand() % (n - len);
if (mp.count(k)) return mp[k];
return k;
}
};
预处理时间复杂度 O ( n − l B ) O(n-l_B) O(n−lB),pick时间 O ( 1 ) O(1) O(1),空间 O ( l B ) O(l_B) O(lB)。