Memory efficient numpy.random.choice()

NumPy’s numpy.random.choice() method samples items at random from an array.

Unfortunately, sampling without replacement an int from to requires memory. Thus, the following code will result in an OOM error (if not completely crash your machine):

import numpy as np

N = 1_000_000_000_000
np.random.choice(N, size=2, replace=False)

This is because numpy.random.choice() creates (at least up to NumPy version 1.24) an array of size (in our example, ~465 GB !), as it can be seen in the source code:

    def choice(self, a, size=None, replace=True, p=None):

        ...

                # pop_size == a when `a` is an `int`
                idx = self.permutation(pop_size)[:size]

        ...

        if a.ndim == 0:
            return idx


    def permutation(self, object x):

        ...

        if isinstance(x, (int, np.integer)):
            # this creates an array of size `x`
            arr = np.arange(x)
            self.shuffle(arr)
            return arr

However, intuitively we shouldn’t need to create an array of size if we are sampling a small amount of items , right?

We could simply keep sampling with replacement (which doesn’t require memory) until we have an array of different items.

This very basic and straightforward algorithm would be defined as follows:

Algorithm 1 An algorithm to draw without replacement. items are taken without replacement from a universe of size . samples items with replacement with a space complexity of :

Require:

1:
2: while do
3:
4: if then
5:
6: end if
7: end while

Which in code looks like,

def choice(N: int, k: int = 1, replace: bool = True, p=None):
    """Sample `k` elements from 0 to `N-1` with or without replacement."""
    if not replace and N >= 3.414213562373095 * k:
        X = [np.random.choice(N, size=1, replace=True, p=p)]
        X = set(X)
        while len(X) != k:
            s = np.random.choice(N, size=1, replace=True, p=p)
            X.add(s)
        return X
    else:
        return np.random.choice(N, size=k, replace=replace, p=p)

But why require that ? We will figure it out by proving a couple of lemmas:

Lemma 1 Algorithm 1 has a space complexity of .

Proof. Both and occupy space, since is at most and our largest number, , requires bits to be stored.

This one was very simple. But, how about time complexity?

Lemma 2 Algorithm 1 has an expected time complexity of , i.e., it has a time complexity of .

Proof. In order to advance in the while loop, we need to add new elements to . Thus:

Using this probability bound, the expected number of items drawn is:

Since we have to draw 1 item in line 1 of the algorithm and then we have to draw new unique items in the while loop.

We now note that the algorithm requires that , so:

Note that sets in Python have lookup and insert times, required for lines 4 and 5 in the algorithm.

However, this is not all. We have used the fact that in our proof, but now let’s generalize to with . Note that would include the possibility that , in which we just sample all items in the universe of size , and this is not interesting for our purposes. Also, includes cases where , which is impossible and our algorithm would not halt.

Lemma 3 If for , the expected number of items drawn is less than .

Proof. implies that,

Since,

With this lemma we can prove the following lemma:

Lemma 4 Algorithm 1 draws O(2k) items w.p. (with probability) less than .

Proof. We define , which is the amount of items drawn in the while loop. We recall Markov’s inequality:

So that, using Lemma 3,

Now, we set , and consider the maximum value of , ,

Note that this is true since Lemma 3 still applies when . We now observe that,

Thus finally,

We can now apply our algorithm requirement so we get,

Thus, in the worst case we always draw items w.p. less than .

This is very useful. It means that if is massive and is minuscule in comparison, we won’t draw more than items with high probability. In fact, if is 100 times , then we will draw items w.p. less than ~1%.

This is definitely better than the original algorithm, which when creating the array of length had to spend time.

To end this little exploration, note that in general if we had chosen for some , then we would have that the algorithm is w.p. less than .

For example, the algorithm is w.p. less than . This means that drawing 10 items without replacement from a universe of 1’000 items using this algorithm will require drawing 103 items w.p. less than 0.1%.

Example use-case

But, when do we have a massive in practice?

Let me present the original problem I was trying to solve which required me to create my own choice function: a memory-efficient implementation to sample without replacement itertools.product().

Let’s say that you have multiple iterables that can be indexed in Python and you want to sample random permutations of their items.

Since the space of all permutations is prohibitively large, we cannot simply generate all permutations first and then sample of them.

The solution is to use our function:

def random_product(*iterables, size: int = None) -> tuple:
    """Memory efficient way to sample ``itertools.product`` at random without replacement."""
    lens = [len(it) for it in iterables]
    # N is usually extremely large!
    N = np.prod(lens)
    if size is None:
        size = N
    idx = np.unravel_index(choice(N, k=size, replace=False),
                           shape=lens)
    idx = np.array(idx).T
    for i_ in idx:
        yield tuple(iterables[i][i_[i]] for i in range(len(iterables)))