从带索引的numpy数组中随机抽取一个样本

Question

我有一个大小为12x12的包含概率的numpy矩阵。其目的是使用这些概率随机抽取样本，然后返回样本的索引。

目前，我正在使用以下基于np.random.choice的代码来实现这一点，其中grid=numpy矩阵：

rnd_choice = np.random.choice(grid.size, p=grid.ravel() / grid.ravel().sum())
sample_index = np.unravel_index(rnd_choice, grid.shape)

问题是速度，因为在整个模拟过程中我必须这样做几千次。Snakeviz强调这是一个需要改进的领域，因此，我想知道是否有人对如何提高速度有任何想法？

上面代码中使用的Python版本是Python 3.8。

Answer 1

以下内容似乎有点快(在我的笔记本电脑上是x4(：

c = grid.ravel().cumsum()                                              
out = np.unravel_index(c.searchsorted(rng.uniform(0,c[-1])),grid.shape)

Answer 2

如果grid.ravel需要在for循环中计算，因为在for循环的每个循环中概率都会发生变化，那么通过每个循环只调用一次.ravel()，您仍然可以潜在地将计算负担减少两倍：

for cycle in loop:
# grid gets renewed here...
g_ravel = grid.ravel()         # do the ravel() process only once
rnd_choice = np.random.choice(grid.size, p=g_ravel / g_ravel.sum())
sample_index = np.unravel_index(rnd_choice, grid.shape)