Problem statement

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

Solution 1

Here we use similar idea as in problem 0519: Random Flip Matrix: imagine, we have set [0, 1, 2, 3, 4, 5] and we delete 0 and 3 from it. Then we want to make correspondence between [0, 1, 2, 3] and [1, 2, 4, 5], we can create connections: {3: 5, 0: 4}: the size of this dictionary will be O(B), where B is size of blacklist.

Complexity

Time complexity of preprocessing is O(B log B), space is O(B). Time complexity of sample is O(1).

Code

class Solution:
    def __init__(self, N: int, blacklist: List[int]):
        self.N, self.dic = N, {}
        for num in sorted(blacklist, reverse = True):
            self.dic[num] = self.dic.get(self.N-1, self.N-1)
            self.N -= 1
        
    def pick(self):
        rand = random.randint(0, self.N-1)
        return self.dic.get(rand, rand)

Solution 2

We can do similar algorithm in O(B) preprocessing time, if we do not sort blacklist and deal only with numbers, which are small than N-B: we try to add element to dictionary if it is not occupied yet and if it is, we decrease N by one.

Complexity

Time complexity for init is O(B), space is O(B). Time complexity of sample is O(1).

Code

class Solution:
    def __init__(self, N, blacklist):
        B, set_bl = len(blacklist), set(blacklist)
        self.M, self.dic = N - B, {}
        for num in blacklist:
            if num < self.M:
                while N-1 in set_bl: N-=1
                self.dic[num] = N-1
                N -= 1     

    def pick(self):
        rand = random.randint(0, self.M - 1)
        return self.dic.get(rand, rand)