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
26
27
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) {
        // Find the slot that haven't been used.
        while (map.count(maxAvailable))
          --maxAvailable;
        map[b] = maxAvailable--;
      }
  }

  int pick() {
    const int num = rand() % validRange;
    const auto it = map.find(num);
    return it == map.cend() ? num : it->second;
  }

 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
27
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) {
        // Find the slot that haven't been used.
        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
23
24
25
class Solution:
  def __init__(self, n: int, blacklist: List[int]):
    self.validRange = n - len(blacklist)
    self.dict = {}

    maxAvailable = n - 1

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

    for b in blacklist:
      if b < self.validRange:
        # Find the slot that haven't been used.
        while maxAvailable in self.dict:
          maxAvailable -= 1
        self.dict[b] = maxAvailable
        maxAvailable -= 1

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

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

    return value