Vektorisierte Möglichkeit, Elemente mithilfe einer Indexkarte hinzuzufügen?

Question

Ich habe zwei numpy Arrays, eine der Größe (386, 3, 4) und eine andere der Größe (386, 4), die ich values bzw. keys nennen werde. Das zweite Array enthält Ganzzahlen, die Indizes für mein Ausgabe-Array sind. Ich brauche die folgenden for Schleife zu implementieren -

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

Natürlich output hat Abmessungen (max_index + 1, 3). Konnte ich einer vektorisierten Implementierung weichen?

Answer 1

denke ich np.add.at sollte das tun, was Sie wollen:

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

Kleine Array Beispiel:

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

Ansatz # 1

Hier ist ein Ansatz mit 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))) 

---

Ansatz # 2

Hier ist eine andere mit np.add.reduceat nach der Säulen basiert das Sortieren auf die 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