消除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]

這些循環通常在1000x1000的范圍內運行。 如何正確使用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())