Index 2D numpy Array von einem 2D-Array von Indizes ohne Loops

Question

Ich bin auf der Suche nach einem vektorisierten Weg, um eine numpy.array von numpy.array von Indizes zu indexieren.

Zum Beispiel:

import numpy as np 

a = np.array([[0,3,4], 
       [5,6,0], 
       [0,1,9]]) 

inds = np.array([[0,1], 
       [1,2], 
       [0,2]]) 

Ich möchte ein neues Array aufzubauen, so dass jede Zeile (i) in diesem Array ist eine Zeile (i) der Array a durch Zeile des Arrays inds indexiert (i). Meine gewünschte Ausgabe ist:

array([[ 0., 3.], # a[0][:,[0,1]] 
     [ 6., 0.], # a[1][:,[1,2]] 
     [ 0., 9.]]) # a[2][:,[0,2]] 

ich dies mit einer Schleife erreichen kann:

def loop_way(my_array, my_indices): 
    new_array = np.empty(my_indices.shape) 
    for i in xrange(len(my_indices)): 
     new_array[i, :] = my_array[i][:, my_indices[i]] 
    return new_array 

Aber ich bin auf der Suche nach einer reinen vektorisiert Lösung.

Answer 1

Bei der Verwendung von Indexarrays zum Indizieren eines anderen Arrays sollte die Form jedes Indexarrays mit der Form des Ausgangs Array übereinstimmen. Sie wollen, dass die Spaltenindizes inds übereinstimmen, und Sie wollen, dass die Zeilenindizes die Zeile der Ausgabe anzupassen, so etwas wie:

array([[0, 0], 
     [1, 1], 
     [2, 2]]) 

können Sie nur eine einzige Spalte der obigen verwenden, aufgrund Rundfunk, so dass Sie np.arange(3)[:,None] ist die vertikalen arange weil None Einsätze eine neue Achse verwenden:

>>> np.arange(3)[:, None] 
array([[0], 
     [1], 
     [2]]) 

schließlich zusammen:

>>> a[np.arange(3)[:,None], inds] 
array([[0, 3], # a[0,[0,1]] 
     [6, 0], # a[1,[1,2]] 
     [0, 9]]) # a[2,[0,2]] 

Answer 2

Es ist möglich, wenn auch etwas nicht offensichtlich, dies zu tun, wie folgt:

>>> a[np.arange(a.shape[0])[:, None], inds] 
array([[0, 3], 
     [6, 0], 
     [0, 9]]) 

Der Index np.arange(a.shape[0]) einfach indexiert die Reihen der Anordnung von Spaltenindizes, auf die inds gilt. Das Anhängen [:, None] modifiziert die Form dieser Anordnung derart, dass ihre Form (a.shape[0], 1) ist, d. H. Jeder Reihenindex ist in einer separaten Reihe einer 1 Spalten breiten 2D Anordnung.

Das Grundprinzip ist, dass die Anzahl der Dimensionen in den Index-Arrays übereinstimmen muss, und ihre Formen müssen dies auch tun. Siehe Dokumentation für np.ix_, um ein Gefühl dafür zu bekommen.