从多维数组矢量化采样[重复]答案

【问题标题】：Vectorize sampling from a multidimensional array [duplicate]从多维数组矢量化采样[重复]
【发布时间】：2021-09-12 20:24:53
【问题描述】：

我有一个形状为 D x N x K 的 numpy 数组。

我需要一个由 K 个类中的随机元素组成的 D x N 数组，其中每个索引 [d, n] 有一个不同的类概率向量，由第三个轴表示。

np.random.choice 的 numpy documentation 只允许概率的一维数组。

我可以对这种类型的采样进行矢量化，还是必须使用 for 循环，如下所示：

# input_array of shape (D, N, K)
# output_array of shape (D, N)

for d in range(input_array.shape[0]):
    for n in range(input_array.shape[1]):
        probabilities = input_array[d, n]
        element = np.random.choice(K, p=probabilities)
        output_array[d, n] = element

如果有这样的功能我会喜欢的

output_array = np.random.choice(input_array, K, probability_axis=-1)

编辑：设法找到“手工设计”的解决方案here。

【问题讨论】：

rng.choice 允许使用 axis 参数，但 p 仍然必须是 1d。因此，对于 nd 数组，这就像在 arange 上进行选择并将其用作索引。
谢谢，但找到了一个手工设计的解决方案。见编辑。
打败我。我仍然认为我的解决方案稍微简单一些。

标签： python numpy random vectorization probability

【解决方案1】：

np.random.choice 和 np.random.default_rng().choice 都不支持以您想要的方式广播概率。但是，您可以使用np.cumsum 拼凑一些类似的东西：

sprob = input_array.cumsum(axis=-1, dtype=float)
sprob /= sprob[:, :, -1:]
output_array = (np.random.rand(D, N, 1) > sprob).argmin(-1)

很遗憾，np.searchsorted 也不支持多维查找（可能出于相关原因）。

【讨论】：