消除在numpy的循環 - 適當的廣播

Question

我有這個醜陋的循環：

import numpy as np 

in1 = np.array(((2.2, 2.5) , 
        (1.4, 1.7)) 
      ) 
out1 = np.array((0.0, 0.0, 0.0)) 
idx1 = np.array(((0, 1), 
        (2, 2)) 
       ) 

# an ugly loop 
for i in xrange(2): 
    for j in xrange(2): 
    out1[ idx1[i,j] ] += in1[i,j] 

這些循環正常運行在1000×1000的範圍。如何正確使用numpy索引？這不起作用：

out1[ idx1 ] += in1 

Answer 1

我們可以扁平化指數和要分配的值，然後用np.bincount得到輸出，而不需要初始化，像這樣 -

out1 = np.bincount(idx1.ravel(), in1.ravel()) 

另外，要添加到初始化輸出數組out1，我們可以用np.add.at，像這樣 -

np.add.at(out1, idx1.ravel(), in1.ravel())