【Leetcode】710. Random Pick with Blacklist

题目地址：

https://leetcode.com/problems/random-pick-with-blacklist/description/

有一个长$n$的数组$A$，再给定一个数组$B$，其表示$A$中的若干个下标的黑名单。要求实现一个方法，该方法获取黑名单下标之外的随机的$0\sim n-1$的下标，每个下标要等概率返回。

设$m=l_B$，那么$A$的前$n-m$个数里如果有$x$个下标在黑名单之中，那么A的后$m$个数一定有$x$个下标不在黑名单之中，我们可以在前$x$个在黑名单之中的下标到后$x$个不在黑名单的下标做一个映射。每次随机一个$[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)$，pick时间$O(1)$，空间$O(l_B)$。