二維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%那裏的運行時間。