python指定概率选择元素_python – 给定iterable,用概率p独立选择每个元素

直到第一次成功的伯努利试验数量由

geometric distribution表示.因此,您可以使用它来生成要使用numpy.random.geometric跳过的项目数：

import itertools

import numpy

def binomial_choice(L, p):

iterator = iter(L)

while True:

to_skip = numpy.random.geometric(p) - 1

yield next(itertools.islice(iterator, to_skip, None))

这适用于任何迭代器,您无需事先知道长度.

但是对于Python 3.5,由于PEP 479,你必须使用更复杂的版本：

def binomial_choice(L, p):

iterator = iter(L)

try:

while True:

to_skip = numpy.random.geometric(p) - 1

yield next(itertools.islice(iterator, to_skip, None))

except StopIteration:

return

用法示例：

In [1]: list(binomial_choice(range(100), 0.05))

Out[1]: [9, 15, 31, 53, 92, 93]

In [2]: list(binomial_choice(range(5), 1))

Out[2]: [0, 1, 2, 3, 4]

分布似乎非常正确：

In [5]: sum(len(lis