基于阵列元的雷达随机指数采样

Question

我有一个ndarray，其中的元素之和为1。我的问题很简单，我想随机抽样这个数组的索引，其中得到每个索引的概率是相应元素的值。现在，我想做很多次，在大数组中，有很多维度，所以我希望使用np.random.choice或类似的东西，而不是for-循环，但我似乎找不出答案。

例如，如果我有

import numpy as np
a = np.array([[0.0, 0.2], [0.6, 0.2]])

我不想以20 %的速度获得索引(0,0)、get、(0,1)和(1,1)，而以60 %的速度获得(1,0)。

据你所知，有什么好的、快捷的方法吗？非常感谢！

Answer 1

您可以使用flatten数组(从2D到1D)并使用np.random.choice。然后，您需要使用np.unravel_index将一维索引投影回2D。

np.unravel_index(np.random.choice(a.size, p = a.flatten()), a.shape)