710. Random Pick with Blacklist

Given a blacklist B containing unique integers from [0, N), write a function to return a uniform random integer from [0, N) which is NOT in B.

Optimize it such that it minimizes the call to system’s Math.random().

Note:

  1. 1 <= N <= 1000000000
  2. 0 <= B.length < min(100000, N)
  3. [0, N) does NOT include N. See interval notation.

Example 1:

Input: 
["Solution","pick","pick","pick"]
[[1,[]],[],[],[]]
Output: [null,0,0,0]

Example 2:

Input: 
["Solution","pick","pick","pick"]
[[2,[]],[],[],[]]
Output: [null,1,1,1]

Example 3:

Input: 
["Solution","pick","pick","pick"]
[[3,[1]],[],[],[]]
Output: [null,0,0,2]

Example 4:

Input: 
["Solution","pick","pick","pick"]
[[4,[2]],[],[],[]]
Output: [null,1,3,1]

Explanation of Input Syntax:

The input is two lists: the subroutines called and their arguments. Solution's constructor has two arguments, N and the blacklist B. pick has no arguments. Arguments are always wrapped with a list, even if there aren't any.

Rust Solution

use rand::prelude::*;
use std::collections::HashMap;
use std::collections::HashSet;

struct Solution {
    rng: ThreadRng,
    map: HashMap<usize, usize>,
    m: usize,
    n: usize,
}

impl Solution {
    fn new(n: i32, mut blacklist: Vec<i32>) -> Self {
        let n = n as usize;
        let rng = thread_rng();
        blacklist.sort_unstable();
        let m = blacklist.len();
        let set: HashSet<i32> = blacklist.iter().copied().collect();
        let mut map: HashMap<usize, usize> = HashMap::new();
        let mut j = 0;
        for i in n - m..n {
            if !set.contains(&(i as i32)) {
                map.insert(blacklist[j] as usize, i);
                j += 1;
            }
        }
        Solution { rng, map, m, n }
    }

    fn pick(&mut self) -> i32 {
        let x = self.rng.gen_range(0, self.n - self.m);
        if let Some(&v) = self.map.get(&x) {
            v as i32
        } else {
            x as i32
        }
    }
}

Having problems with this solution? Click here to submit an issue on github.