How to efficiently use numpy random choice for varying weight list

user_11958218 0
Viewed 254

I am trying to use the function np.random.choice to randomly choose numbers from a list whose weights are in a list of lists. That is, for every row I want to generate one number. It is possible to do it with for loop as follows,

from numpy.random import choice

W_list = np.array([[0.9,0.1],
                   [0.95,0.05],
                   [0.85,0.15]])
number_list = []
for i in range(len(W_list)):
    number_list.extend(choice([0, 1], size=1, p=W_list[i]).tolist())

number_list

[0,0,0]

The p parameter needs to 1D, hence it is not possible to use p=W_list. Is there any way to do this more efficiently without using the for loop?

random numpy python
1 Answers

For the simple case of a single boolean per row, you can do this very easily by implementing the way probabilities are applied by hand:

cumulative_weights = W_list.cumsum(axis=1)
number_arr = (np.random.rand(3, 1) < cumulative_weights).argmax(axis=1)
user_29887300
Related