LeetCode-in-Java

710. Random Pick with Blacklist

Hard

You are given an integer n and an array of unique integers blacklist. Design an algorithm to pick a random integer in the range [0, n - 1] that is not in blacklist. Any integer that is in the mentioned range and not in blacklist should be equally likely to be returned.

Optimize your algorithm such that it minimizes the number of calls to the built-in random function of your language.

Implement the Solution class:

Example 1:

Input

[“Solution”, “pick”, “pick”, “pick”, “pick”, “pick”, “pick”, “pick”]

[[7, [2, 3, 5]], [], [], [], [], [], [], []]

Output: [null, 0, 4, 1, 6, 1, 0, 4]

Explanation:

Solution solution = new Solution(7, [2, 3, 5]);
solution.pick(); // return 0, any integer from [0,1,4,6] should be ok. Note that for every call of pick,
                 // 0, 1, 4, and 6 must be equally likely to be returned (i.e., with probability 1/4).
solution.pick(); // return 4
solution.pick(); // return 1
solution.pick(); // return 6
solution.pick(); // return 1
solution.pick(); // return 0
solution.pick(); // return 4 

Constraints:

Solution

import java.util.HashMap;
import java.util.Map;
import java.util.Random;

@SuppressWarnings("java:S2245")
public class Solution {
    private final Map<Integer, Integer> map;
    private final Random r;
    private final int upperLimit;

    public Solution(int n, int[] blacklist) {
        map = new HashMap<>();
        r = new Random();
        upperLimit = n - blacklist.length;
        for (int val : blacklist) {
            map.put(val, -1);
        }
        int j = n - 1;
        for (int val : blacklist) {
            if (val < upperLimit) {
                while (map.containsKey(j)) {
                    j--;
                }
                map.put(val, j);
                j--;
            }
        }
    }

    public int pick() {
        int val = r.nextInt(upperLimit);
        if (map.containsKey(val)) {
            return map.get(val);
        }
        return val;
    }
}