Python：从 n 维数组中定义的离散分布采样答案

【问题标题】：Python: Sampling from a discrete distribution defined in an n-dimensional arrayPython：从 n 维数组中定义的离散分布采样
【发布时间】：2014-08-26 05:30:37
【问题描述】：

Python 中是否有一个函数可以从 n 维 numpy 数组中采样并返回每次绘制的索引。如果不是，如何定义这样一个函数？

例如：

>>> probabilities = np.array([[.1, .2, .1], [.05, .5, .05]])  
>>> print function(probabilities, draws = 10)
 ([1,1],[0,2],[1,1],[1,0],[0,1],[0,1],[1,1],[0,0],[1,1],[0,1])

我知道这个问题可以通过一维数组以多种方式解决。但是，我将处理大型 n 维数组，不能仅仅为了一次绘制而对它们进行整形。

【问题讨论】：

标签： python arrays numpy random-sample

【解决方案1】：

你可以使用np.unravel_index:

a = np.random.rand(3, 4, 5)
a /= a.sum()

def sample(a, n=1):
    a = np.asarray(a)
    choices = np.prod(a.shape)
    index = np.random.choice(choices, size=n, p=a.ravel())
    return np.unravel_index(index, dims=a.shape)

>>> sample(a, 4)
(array([2, 2, 0, 2]), array([0, 1, 3, 2]), array([2, 4, 2, 1]))

这将返回一个数组元组，每个维度为a，每个长度为请求的样本数。如果您希望有一个形状为(samples, dimensions) 的数组，请将return 语句更改为：

return np.column_stack(np.unravel_index(index, dims=a.shape))

现在：

>>> sample(a, 4)
array([[2, 0, 0],
       [2, 2, 4],
       [2, 0, 0],
       [1, 0, 4]])

【讨论】：

【解决方案2】：

如果您的数组在内存中是连续的，您可以更改数组的shape：

probabilities = np.array([[.1, .2, .1], [.05, .5, .05]]) 
nrow, ncol = probabilities.shape
idx = np.arange( nrow * ncol ) # create 1D index

probabilities.shape = ( 6, ) # this is OK because your array is contiguous in memory

samples = np.random.choice( idx, 10, p=probabilities ) # sample in 1D
rowIndex = samples / nrow # convert to 2D
colIndex = samples % ncol

array([2, 0, 1, 0, 2, 2, 2, 2, 2, 0])
array([1, 1, 2, 0, 1, 1, 1, 1, 1, 1])

请注意，由于您的数组在内存中是连续的，reshape 也会返回一个视图：

In [53]:

view = probabilities.reshape( 6, -1 )
view[ 0 ] = 9
probabilities[ 0, 0 ]
Out[53]:
9.0

【讨论】：

谢谢，这将如何推广到 n 维？