从瓮中绘制的 Numpy

Question

我想在numpy中运行一个比较简单的随机抽奖，但是找不到好的表达方式。 我认为最好的方法是将其描述为从骨灰盒中绘制而无需更换。我有一个有 k 种颜色的瓮，以及每种颜色的 n_k 个球。我想画m个球，并且知道我有多少个每种颜色的球。

我目前的尝试

np.bincount(np.random.permutation(np.repeat(np.arange(k), n_k))[:m], minlength=k)

这里，n_k 是一个长度为 k 的数组，其中包含球的数量。

似乎相当于 np.bincount(np.random.choice(k, m, n_k / n_k.sum(), minlength=k)

这有点好，但仍然不是很好。

Answer 1

您想要的是multivariate hypergeometric distribution 的实现。 我不知道 numpy 或 scipy 中的哪一个，但它可能已经存在于某个地方。

我为 numpy 1.18.0 贡献了多元超几何分布的实现；见numpy.random.Generator.multivariate_hypergeometric。

例如，从一个装有 12 个红色、4 个绿色和 18 个蓝色弹珠的瓮中抽取 15 个样本，并重复该过程 10 次：

In [4]: import numpy as np

In [5]: rng = np.random.default_rng()

In [6]: colors = [12, 4, 18]

In [7]: rng.multivariate_hypergeometric(colors, 15, size=10)                    
Out[7]: 
array([[ 5,  4,  6],
       [ 3,  3,  9],
       [ 6,  2,  7],
       [ 7,  2,  6],
       [ 3,  0, 12],
       [ 5,  2,  8],
       [ 6,  2,  7],
       [ 7,  1,  7],
       [ 8,  1,  6],
       [ 6,  1,  8]])

这个答案的其余部分现在已经过时了，但我会留给后代（不管这意味着什么......）。

---

您可以通过重复调用numpy.random.hypergeometric 来实现它。这是否比您的实施更有效取决于有多少种颜色以及每种颜色有多少个球。

例如，下面的脚本打印从包含三种颜色（红色、绿色和蓝色）的瓮中绘制的结果：

from __future__ import print_function

import numpy as np


nred = 12
ngreen = 4
nblue = 18

m = 15

red = np.random.hypergeometric(nred, ngreen + nblue, m)
green = np.random.hypergeometric(ngreen, nblue, m - red)
blue = m - (red + green)

print("red:   %2i" % red)
print("green: %2i" % green)
print("blue:  %2i" % blue)

样本输出：

red:    6
green:  1
blue:   8

以下函数将其概括为在给定数组colors 保存每种颜色的数量的情况下选择m 球：

def sample(m, colors):
    """
    Parameters
    ----------
    m : number balls to draw from the urn
    colors : one-dimensional array of number balls of each color in the urn

    Returns
    -------
    One-dimensional array with the same length as `colors` containing the
    number of balls of each color in a random sample.
    """

    remaining = np.cumsum(colors[::-1])[::-1]
    result = np.zeros(len(colors), dtype=np.int)
    for i in range(len(colors)-1):
        if m < 1:
            break
        result[i] = np.random.hypergeometric(colors[i], remaining[i+1], m)
        m -= result[i]
    result[-1] = m
    return result

例如，

>>> sample(10, [2, 4, 8, 16])
array([2, 3, 1, 4])

Answer 2

以下应该有效：

def make_sampling_arr(n_k):
    out = [ x for s in [ [i] * n_k[i] for i in range(len(n_k)) ] for x in s ]
    return out

np.random.choice(make_sampling_arr(n_k), m)