python - 给定 iterable，以概率 p 独立选择每个元素

Question

我需要一个函数来在 Python 中执行以下操作。

我需要遍历输入列表并以概率 p(很小)选择每个元素。

这是一个简单的实现，可以清楚地说明我在做什么。

def f(inputList,p):
    for element in inputList:
        if random.random()<p:
            yield element

每次随机数生成都是昂贵的。我们可以通过最初计算第一次成功之前需要多长时间然后跳转到该条目来减少随机数生成。我有一个方法，但想知道是否已经存在或更好的编码方法。原则上我只需要它用于列表，但我想要一些适用于一般可迭代对象的东西。

def jump_calculation(p):
    if p == 1:
        return 0
    r = random.random()
    k = int(scipy.log(1-r)/scipy.log(1-p))
    return k
def binomial_choice(L,p):
    jump = jump_calculation(p)
    for element in L:
        jump -= 1
        if jump<0:
            yield element
            jump = jump_calculation(p)

我是在重新发明轮子吗？如果不是，是否有明显的改进使代码更易于阅读？

Answer 1

第一次成功之前的伯努利试验次数由 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(list(binomial_choice(range(100), 0.05))) for i in range(100000)) / 100000
Out[5]: 4.99883

它也比您的两种方法更快:

In [14]: timeit list(binomial_choice_geometric(range(1000), 0.01))
10000 loops, best of 3: 24.4 µs per loop

In [11]: timeit list(binomial_choice_geometric_3_5(range(1000), 0.01))
10000 loops, best of 3: 42.7 µs per loop

In [12]: timeit list(binomial_choice_jump_calculation(range(1000), 0.01))
1000 loops, best of 3: 596 µs per loop

In [13]: timeit list(binomial_choice_foreach_random(range(1000), 0.01))
1000 loops, best of 3: 203 µs per loop

它实际上在另一个答案的 random.sample 方法的规模上运行(根据评论中的建议修改为使用 numpy.random.binomial 来获得正确的分布)，但不需要有一个列表，也不需要提前知道参数的 len:

In [19]: timeit list(binomial_choice_random_sample(range(1000), 0.01))
10000 loops, best of 3: 19.8 µs per loop