来自2d概率numpy数组的样本？

Question

说我有一个二维数组ar，像这样：

0.9, 0.1, 0.3
0.4, 0.5, 0.1
0.5, 0.8, 0.5

我想根据此概率数组从[1，0]进行采样。

rdchoice = lambda x: numpy.random.choice([1, 0], p=[x, 1-x])

我尝试了两种方法：

1）首先将其重塑为1d数组，然后使用numpy.random.choice，然后将其重塑为2d：

np.array(list(map(rdchoice, ar.reshape((-1,))))).reshape(ar.shape)

2）使用向量化功能。

func = numpy.vectorize(rdchoice)
func(ar)

但是这两种方法都太慢了，我了解到矢量化的本质是一个for循环，在我的实验中，我发现map并不比vectorize快。

我认为这可以更快地完成。如果二维数组很大，那就太慢了。

Answer 1

您应该能够这样做：

>>> p = np.array([[0.9, 0.1, 0.3], [0.4, 0.5, 0.1], [0.5, 0.8, 0.5]])
>>> (np.random.rand(*p.shape) < p).astype(int)

Answer 2

实际上我可以使用np.random.binomial：

import numpy as np
p = [[0.9, 0.1, 0.3],
     [0.4, 0.5, 0.1],
     [0.5, 0.8, 0.5]]

np.random.binomial(1, p)