二维numpy数组的加权随机采样

Question

我有一个2d numpy数组Z，我想随机选择一个索引Z，其中索引被选择的机会与该索引处的Z值成正比。

现在，我做了以下内容：

yar = list(np.ndenumerate(Z)) 
x,y = yar[np.random.choice(len(yar), p=Z.ravel()/Z.sum())][0] 

哪个做这项工作，但感觉丑陋的（并且是极其缓慢之外）。有没有更好的办法？

Answer 1

我们可以优化避免创建yar。我们只需从np.random.choice获得相应的线性指标，将其转换为维数指数np.unravel_index即可得到x和y。

因此，实现起来 -

linear_idx = np.random.choice(Z.size, p=Z.ravel()/float(Z.sum())) 
x, y = np.unravel_index(linear_idx, Z.shape) 

只给对其中yar创作是造成瓶颈在设置的数字有些情况下，这里的取样时间测试 -

In [402]: Z = np.random.randint(0,9,(300,400)) 

In [403]: yar = list(np.ndenumerate(Z)) 

In [404]: %timeit list(np.ndenumerate(Z)) 
10 loops, best of 3: 46.3 ms per loop 

In [405]: %timeit yar[np.random.choice(len(yar), p=Z.ravel()/float(Z.sum()))][0] 
1000 loops, best of 3: 1.34 ms per loop 

In [406]: 46.3/(46.3+1.34) 
Out[406]: 0.971872376154492 

因此，创建yar是吞吃97%那里的运行时间。