以不同概率选择列表元素的Pythonic方法

Question

import random
pos = ["A", "B", "C"]
x = random.choice["A", "B", "C"]

这段代码让我以相等的概率得到"A"、"B"或"C"。如果我想要"A"的概率是30%，"B"的概率是40%，而"C"的概率是30%，有没有什么好的方法来表示呢？

Answer 1

这里有一个类，可以展示一堆物品及其相对概率，而不需要实际列出所有的物品：

import bisect
class WeightedTuple(object):
    """
    >>> p = WeightedTuple({'A': 2, 'B': 1, 'C': 3})
    >>> len(p)
    6
    >>> p[0], p[1], p[2], p[3], p[4], p[5]
    ('A', 'A', 'B', 'C', 'C', 'C')
    >>> p[-1], p[-2], p[-3], p[-4], p[-5], p[-6]
    ('C', 'C', 'C', 'B', 'A', 'A')
    >>> p[6]
    Traceback (most recent call last):
    ...
    IndexError
    >>> p[-7]
    Traceback (most recent call last):
    ...
    IndexError
    """
    def __init__(self, items):
        self.indexes = []
        self.items = []
        next_index = 0
        for key in sorted(items.keys()):
            val = items[key]
            self.indexes.append(next_index)
            self.items.append(key)
            next_index += val

        self.len = next_index

    def __getitem__(self, n):
        if n < 0:
            n = self.len + n
        if n < 0 or n >= self.len:
            raise IndexError

        idx = bisect.bisect_right(self.indexes, n)
        return self.items[idx-1]

    def __len__(self):
        return self.len

现在，只需这样说：

data = WeightedTuple({'A': 30, 'B': 40, 'C': 30})
random.choice(data)

Answer 2

不是……那么多……

pos = ['A'] * 3 + ['B'] * 4 + ['C'] * 3
print random.choice(pos)

或者

pos = {'A': 3, 'B': 4, 'C': 3}
print random.choice([x for x in pos for y in range(pos[x])])

Answer 3

权重定义了一种概率分布函数（pdf）。我们可以通过将均匀随机数（在0到1之间）代入与之相关的反累积分布函数，来生成这种分布下的随机数。

你也可以看看这个StackOverflow的解释，或者像维基百科上所说的：

如果Y的分布是U[0,1]，那么F⁻¹(Y)的分布就是F。这在使用反变换抽样方法生成随机数时会用到。

import random
import bisect
import collections

def cdf(weights):
    total = sum(weights)
    result = []
    cumsum = 0
    for w in weights:
        cumsum += w
        result.append(cumsum / total)
    return result

def choice(population, weights):
    assert len(population) == len(weights)
    cdf_vals = cdf(weights)
    x = random.random()
    idx = bisect.bisect(cdf_vals, x)
    return population[idx]

weights=[0.3, 0.4, 0.3]
population = 'ABC'
counts = collections.defaultdict(int)
for i in range(10000):
    counts[choice(population, weights)] += 1
print(counts)

# % test.py
# defaultdict(<type 'int'>, {'A': 3066, 'C': 2964, 'B': 3970})

上面的choice函数使用了bisect.bisect，所以选择一个加权随机变量的时间复杂度是O(log n)，其中n是weights的长度。

---

注意，从版本1.7.0开始，NumPy有了一个经过Cython优化的np.random.choice函数。例如，这个函数可以从种群[0,1,2,3]中生成1000个样本，权重为[0.1, 0.2, 0.3, 0.4]：

import numpy as np
np.random.choice(4, 1000, p=[0.1, 0.2, 0.3, 0.4])

np.random.choice还有一个replace参数，可以选择是否允许重复抽样。

---

一个理论上更好的算法是别名法。它会建立一个表格，这个过程需要O(n)的时间，但之后抽样的时间复杂度是O(1)。所以，如果你需要抽取很多样本，理论上别名法可能会更快。这里有一个Walker别名法的Python实现链接，还有一个NumPy版本的链接。