Mapping über 2 numpy.darray gleichzeitig

Question

Hier ist das Problem. Lassen Sie uns sagen, ich habe eine Matrix A =

array([[ 1., 0., 2.], 
     [ 0., 0., 2.], 
     [ 0., -1., 3.]]) 

und einen Vektor der Indizes p = array([0, 2, 1]). Ich möchte eine 3x3-Matrix A in ein Array der Länge 3 (nenne es v) umwandeln, wobei v [j] = A [j, p [j]] für j = 0, 1, 2. Ich kann es auf die folgende Weise tun :

v = map(lambda (row, idx): row[idx], zip(A, p)) 

So für die obige Matrix A und einen Vektor der Indizes p erwarte ich array([1, 2, -1]) (dh 0-ten Element der Reihe 0, 2. Element der Reihe 1, 1. Element der Zeile 2) zu erhalten.

Aber kann ich das gleiche Ergebnis erzielen, indem ich native numpy verwende (dh ohne explizites Zippen und dann Mapping)? Vielen Dank.

Answer 1

Ich glaube nicht, dass eine solche Funktionalität existiert. Um zu erreichen, was Sie wollen, kann ich mir zwei einfache Wege vorstellen. Sie tun können:

np.diag(A[:, p]) 

Hier wird das Array p für jede Zeile als Spaltenindex angewendet wird, so dass auf der Diagonalen werden Sie die Elemente, die Sie für suchen.

Als Alternative können Sie vermeiden, viele unnötige Einträge zu erzeugen, indem mit:

A[np.arange(A.shape[0]), p]