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

Question

假设我有一个像这样的2d数组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）使用vectorize函数。

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

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

我认为这可以更快地完成。如果2d阵列很大，那将是无法忍受的缓慢。

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)