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())

消除numpy中的循环 - 正确的广播

1 个答案: