LeetCode 710: Random Pick with Blacklist

Problem Restatement

We need to design a class that randomly returns an integer from the range:

[0, n - 1]

Some values are blacklisted and cannot be returned.

Every allowed value must have the same probability.

The class supports:

The challenge is to make pick() efficient.

Input and Output

The class shape is:

class Solution:

    def __init__(self, n: int, blacklist: list[int]):
        ...

    def pick(self) -> int:
        ...

Examples

Example:

n = 7
blacklist = [2, 3, 5]

The allowed values are:

[0, 1, 4, 6]

Each call to pick() must return one of these values with equal probability:

1 / 4

Possible outputs:

0
1
4
6

Values 2, 3, and 5 must never appear.

First Thought: Retry Until Valid

One simple approach is:

1. Generate a random number in [0, n - 1].
2. If it is blacklisted, generate another random number.
3. Repeat until we get a valid number.

This works, but it becomes slow when many values are blacklisted.

For example:

n = 1_000_000

and almost all numbers are blacklisted.

Then most random picks are rejected.

Problem With Rejection Sampling

Suppose:

n = 10
blacklist size = 9

Only one value is valid.

Almost every random attempt fails.

The expected runtime becomes poor.

The problem specifically asks us to minimize calls to the random generator.

We need a method that always succeeds in one random pick.

Key Insight

Suppose:

m = len(blacklist)

Then the number of valid values is:

size = n - m

We can generate a random integer:

x in [0, size - 1]

This interval contains exactly as many numbers as there are valid values.

The problem is that some numbers inside this interval may themselves be blacklisted.

For example:

n = 7
blacklist = [2, 3, 5]
size = 4

We randomly choose from:

[0, 1, 2, 3]

But 2 and 3 are blacklisted.

We solve this by remapping.

The large values:

[4, 5, 6]

contain some valid numbers not used in the small range.

Specifically:

4
6

are valid.

So we map:

2 -> 4
3 -> 6

Now every random number from [0, 3] becomes valid:

Each valid number still has equal probability.

Algorithm

Let:

size = n - len(blacklist)

We will only generate random numbers in:

[0, size - 1]

Create:

black = set(blacklist)
mapping = {}

Now find every blacklisted value inside the small range:

b < size

These values need remapping.

We remap them to valid values in the large range:

[size, n - 1]

that are not blacklisted.

During pick():

1. Generate random integer x in [0, size - 1].
2. If x exists in mapping, return mapping[x].
3. Otherwise return x.

Correctness

Let:

size = n - len(blacklist)

There are exactly size allowed values.

The algorithm generates a uniformly random integer x from:

[0, size - 1]

So each x has probability:

1 / size

If x is not blacklisted, the algorithm returns x directly.

If x is blacklisted, then x is remapped to a unique allowed value in the upper range. Each remapped target is chosen so that:

1. It is not blacklisted.
2. It is not used by another mapping.

Therefore every generated x corresponds to exactly one allowed value, and every allowed value corresponds to exactly one generated x.

Since the random generator chooses all x values uniformly, all allowed values are returned with equal probability.

Blacklisted values are never returned because every blacklisted value in the small range is replaced by a valid mapping, and every value in the upper range used for mapping is guaranteed not to be blacklisted.

Thus pick() returns every allowed value uniformly and never returns a blacklisted value.

Complexity

Let:

b = len(blacklist)

Implementation

import random

class Solution:

    def __init__(self, n: int, blacklist: list[int]):
        self.size = n - len(blacklist)

        black = set(blacklist)
        self.mapping = {}

        available = self.size

        for b in blacklist:
            if b < self.size:
                while available in black:
                    available += 1

                self.mapping[b] = available
                available += 1

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

        if x in self.mapping:
            return self.mapping[x]

        return x

Code Explanation

The valid output count is:

self.size = n - len(blacklist)

We only generate random values in:

[0, self.size - 1]

The blacklist is converted to a set for fast lookup:

black = set(blacklist)

We store remappings in:

self.mapping = {}

The variable:

available = self.size

tracks the next valid value in the upper range.

For every blacklisted value inside the small range:

if b < self.size:

we search for a valid upper-range number:

while available in black:
    available += 1

Then create the remapping:

self.mapping[b] = available

During pick():

x = random.randint(0, self.size - 1)

If x needs remapping, return the mapped value:

if x in self.mapping:
    return self.mapping[x]

Otherwise x is already valid:

return x

Example Walkthrough

Suppose:

n = 7
blacklist = [2, 3, 5]

Then:

size = 4

The random range is:

[0, 1, 2, 3]

Blacklisted values inside this range are:

2
3

Upper range values are:

[4, 5, 6]

Ignoring blacklisted 5, we map:

2 -> 4
3 -> 6

Now every random pick becomes valid.

Testing

def test_random_pick_with_blacklist():
    s = Solution(7, [2, 3, 5])

    allowed = {0, 1, 4, 6}

    for _ in range(1000):
        value = s.pick()
        assert value in allowed

    s = Solution(4, [2])

    allowed = {0, 1, 3}

    for _ in range(1000):
        value = s.pick()
        assert value in allowed

    s = Solution(1, [])

    for _ in range(100):
        assert s.pick() == 0

    print("all tests passed")

test_random_pick_with_blacklist()

Test coverage: