Question

我有两个numpy数组，一个大小为(386, 3, 4)，另一个大小为(386, 4)，我将分别称为values和keys。第二个数组包含整数，它是我的输出数组的索引。我需要实现以下for循环 -

for i in range(386):
    for j in range(4):
        output[keys[i, j]] += values[i, :, j]

当然，output的尺寸为(max_index + 1, 3)。我可以使用矢量化实现吗？

Answer 1

我认为np.add.at应该做你想做的事：

np.add.at(output, keys, np.transpose(values, (0, 2, 1)))

小数组示例：

values
# array([[[100, 200, 300, 400],
    [ 10,  20,  30,  40],
    [  1,   2,   3,   4]],

   [[500, 600, 700, 800],
    [ 50,  60,  70,  80],
    [  5,   6,   7,   8]]])
keys
# array([[4, 0, 3, 1],
   [1, 0, 2, 2]])
out
# array([[0, 0, 0],
   [0, 0, 0],
   [0, 0, 0],
   [0, 0, 0],
   [0, 0, 0]])
np.add.at(out, keys, np.transpose(values, (0, 2, 1)))
out
# array([[ 800,   80,    8],
   [ 900,   90,    9],
   [1500,  150,   15],
   [ 300,   30,    3],
   [ 100,   10,    1]])

Answer 2

方法＃1

以下是使用np.tensordot -

的一种方法

# Store size param         
n = values.shape[0]

# Get mask for mapping each key to corresponding row in o/p array
# Simply put : mask = keys==np.arange(n)[:,None,None]        
r,c = np.indices(keys.shape)
mask = np.zeros((keys.max()+1,n,keys.shape[1]),dtype=bool)
mask[keys,r,c] = 1

# Finally mask and sum reduce elems off values
out = np.tensordot(mask, values, axes=((1,2),(0,2)))

方法＃2

根据keys

对列进行排序后，这是另一个np.add.reduceat

n,nr = values.shape[:2]        
kr = keys.ravel()
sidx = kr.argsort()
krs = kr[sidx]
v = values.transpose(1,0,2).reshape(nr,-1)[:,sidx]

cut_idx = np.r_[0,np.flatnonzero(krs[1:] != krs[:-1])+1]
out = np.zeros((keys.max()+1,nr))
out[krs[cut_idx]] = np.add.reduceat(v, cut_idx, axis=1).T

使用索引图添加元素的矢量化方式？

2 个答案: