CODEWITHSUNDEEP
pythonpython-3.xrandomprobability
mel
mel

Reputation: 2790

python: choose element from a list following a certain probability

I'm looking for a python function that would allow me, given a list and a certain probability, to return only the elements that "pass" the lottery. For instance:

my_list = ['A', 'B', 'C']
my_probability = 0.66 # The probability for an element to stay in the list
winner = THE_RANDOM_FUNCTION(my_list, my_probability)
print(winner)
>> ['A', 'C'] # Could be a different result but that only an example

So I want the function to do a lottery (with the given probability) for every element in the list and to keep it in the result list if it passes the lottery.

Is there such function in the standard library or should I implement it?

Upvotes: 1

Views: 1773

Answers (2)

Eli Korvigo
Eli Korvigo

Reputation: 10483

This looks like a classic binomial trial to me. 

In [1]: import numpy as np

In [2]: my_list = ['A', 'B', 'C']

In [3]: prob = 0.66

In [4]: mask = np.random.binomial(1, prob, len(my_list))

In [5]: selected = [elem for keep, elem in zip(mask, my_list) if keep]

In [6]: selected
Out[6]: ['A', 'B']

Or you can go the pure Numpy-way, if you don't need the final result to remain a list

In [7]: my_arr = np.array(my_list)

In [8]: my_arr[mask.astype(bool)]
Out[8]: 
array(['A', 'B'], dtype='<U1')

Upvotes: 1

Tom Zych
Tom Zych

Reputation: 13576

It sounds like you mean you want to choose or not choose each element, with the given probability. Here are a couple of simple solutions:

filter(lambda x: random.random() < my_probability, my_list)

[x for x in my_list if random.random() < my_probability]

Upvotes: 6

Related Questions