インデックスマップを使用してエレメントを追加するベクター化された方法はありますか？

Question

私は2つのnumpyアレイを持っています。サイズは(386, 3, 4)、サイズは(386, 4)です。それぞれvaluesとkeysと呼びます。 2番目の配列には、出力配列のインデックスである整数が含まれています。もちろん

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

、outputは寸法(max_index + 1, 3)を持っている - 私は、次のforループを実装する必要があります。私はベクトル化された実装で道を譲ることができますか？

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