Search code examples
pythonarrayspython-3.xnumpynumpy-ndarray

How to make an numpy array of 0 and 1 based on a probability array?

I know that using Python's random.choices I can do this:

import random


array_probabilities = [0.5 for _ in range(4)]
print(array_probabilities)  # [0.5, 0.5, 0.5, 0.5]

a = [random.choices([0, 1], weights=[1 - probability, probability])[0] for probability in array_probabilities]
print(a)  # [1, 1, 1, 0]

How to make an numpy array of 0 and 1 based on a probability array?

Using random.choices is fast, but I know numpy is even faster. I would like to know how to write the same code but using numpy. I'm just getting started with numpy and would appreciate your feedback.

Solution

  • One option:

    out = (np.random.random(size=len(array_probabilities)) > array_probabilities).astype(int)

    Example output:

    array([0, 1, 0, 1])