Façon vectorisée d'ajouter des éléments en utilisant une carte d'index?

Question

J'ai deux tableaux numpy, un de la taille (386, 3, 4) et un autre de la taille (386, 4), que j'appellerai respectivement values et keys. Le second tableau contient des entiers qui sont des indices pour mon tableau de sortie. Je dois mettre en œuvre la boucle suivante for -

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

Bien sûr, output a des dimensions (max_index + 1, 3). Pourrais-je faire place à une implémentation vectorisée?

Answer 1

Je pense que np.add.at devrait faire ce que vous voulez:

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

exemple Petit tableau:

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

Approche # 1

Voici une approche utilisant 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))) 

---

Approche # 2

Voici une autre avec np.add.reduceat après le tri des colonnes en fonction de la keys -

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