Skip to content

710. Random Pick with Blacklist 👍

  • Time: $O(|\texttt{blacklist}|)$
  • Space: $O(|\texttt{blacklist}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
class Solution {
 public:
  Solution(int N, vector<int>& blacklist) : validRange(N - blacklist.size()) {
    for (const int b : blacklist)
      map[b] = -1;

    int maxAvailable = N - 1;

    for (const int b : blacklist)
      if (b < validRange) {
        while (map.count(maxAvailable))  // find the slot that haven't been used
          --maxAvailable;
        map[b] = maxAvailable--;
      }
  }

  int pick() {
    const int num = rand() % validRange;
    return map.count(num) ? map[num] : num;
  }

 private:
  const int validRange;
  unordered_map<int, int> map;
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
  public Solution(int N, int[] blacklist) {
    validRange = N - blacklist.length;

    for (final int b : blacklist)
      map.put(b, -1);

    int maxAvailable = N - 1;

    for (final int b : blacklist)
      if (b < validRange) {
        while (map.containsKey(maxAvailable))
          --maxAvailable;
        map.put(b, maxAvailable--);
      }
  }

  public int pick() {
    final int num = rand.nextInt(validRange);
    return map.getOrDefault(num, num);
  }

  private int validRange;
  private Map<Integer, Integer> map = new HashMap<>();
  private Random rand = new Random();
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
  def __init__(self, N: int, blacklist: List[int]):
    self.validRange = N - len(blacklist)
    self.dict = {}

    for b in blacklist:
      self.dict[b] = -1

    for b in blacklist:
      if b < self.validRange:
        while N - 1 in self.dict:
          N -= 1
        self.dict[b] = N - 1
        N -= 1

  def pick(self) -> int:
    value = randint(0, self.validRange - 1)

    if value in self.dict:
      return self.dict[value]

    return value