python - 消除 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())