Auswählen bestimmter Zeilen und Spalten aus dem NumPy-Array

Ich bin verrückt geworden, um herauszufinden, was für eine dumme Sache ich hier falsch mache.

Ich verwende NumPy und habe bestimmte Zeilenindizes und bestimmte Spaltenindizes, aus denen ich auswählen möchte. Hier ist der Kern meines Problems:

import numpy as np

a = np.arange(20).reshape((5,4))
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [ 8,  9, 10, 11],
#        [12, 13, 14, 15],
#        [16, 17, 18, 19]])

# If I select certain rows, it works
print a[[0, 1, 3], :]
# array([[ 0,  1,  2,  3],
#        [ 4,  5,  6,  7],
#        [12, 13, 14, 15]])

# If I select certain rows and a single column, it works
print a[[0, 1, 3], 2]
# array([ 2,  6, 14])

# But if I select certain rows AND certain columns, it fails
print a[[0,1,3], [0,2]]
# Traceback (most recent call last):
#   File "<stdin>", line 1, in <module>
# ValueError: shape mismatch: objects cannot be broadcast to a single shape

Warum passiert dies? Sicherlich sollte ich in der Lage sein, die 1., 2. und 4. Zeile sowie die 1. und 3. Spalte auszuwählen? Das erwartete Ergebnis ist:

a[[0,1,3], [0,2]] => [[0,  2],
                      [4,  6],
                      [12, 14]]

Für eine ausgefallene Indizierung müssen Sie alle Indizes für jede Dimension angeben. Sie geben 3 Indizes für den ersten und nur 2 für den zweiten an, daher der Fehler. Sie möchten so etwas tun:

>>> a[[[0, 0], [1, 1], [3, 3]], [[0,2], [0,2], [0, 2]]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Das ist natürlich ein Problem beim Schreiben, sodass Sie sich vom Rundfunk helfen lassen können:

>>> a[[[0], [1], [3]], [0, 2]]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Dies ist viel einfacher, wenn Sie mit Arrays indizieren, nicht mit Listen:

>>> row_idx = np.array([0, 1, 3])
>>> col_idx = np.array([0, 2])
>>> a[row_idx[:, None], col_idx]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Wie Toan vermuten läßt, wäre ein einfacher Hack nur die Zeilen zuerst wählen und dann die Spalten über wählen , dass .

>>> a[[0,1,3], :]            # Returns the rows you want
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [12, 13, 14, 15]])
>>> a[[0,1,3], :][:, [0,2]]  # Selects the columns you want as well
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

[Bearbeiten] Die integrierte Methode: np.ix_

Ich habe kürzlich entdeckt, dass Numpy Ihnen einen eingebauten Einzeiler bietet, mit dem Sie genau das tun können , was @Jaime vorgeschlagen hat, ohne jedoch die Broadcast-Syntax verwenden zu müssen (die unter mangelnder Lesbarkeit leidet). Aus den Dokumenten:

Mit ix_ können schnell Index-Arrays erstellt werden, die das Kreuzprodukt indizieren. a[np.ix_([1,3],[2,5])]Gibt das Array zurück [[a[1,2] a[1,5]], [a[3,2] a[3,5]]].

Also benutzt du es so:

>>> a = np.arange(20).reshape((5,4))
>>> a[np.ix_([0,1,3], [0,2])]
array([[ 0,  2],
       [ 4,  6],
       [12, 14]])

Und so funktioniert es, dass Arrays so ausgerichtet werden, wie Jaime es vorgeschlagen hat, damit die Übertragung ordnungsgemäß erfolgt:

>>> np.ix_([0,1,3], [0,2])
(array([[0],
        [1],
        [3]]), array([[0, 2]]))

Wie MikeC in einem Kommentar sagt, np.ix_hat es auch den Vorteil, eine Ansicht zurückzugeben, was meine erste (vorbearbeitete) Antwort nicht tat. Dies bedeutet, dass Sie jetzt dem indizierten Array zuweisen können :

>>> a[np.ix_([0,1,3], [0,2])] = -1
>>> a    
array([[-1,  1, -1,  3],
       [-1,  5, -1,  7],
       [ 8,  9, 10, 11],
       [-1, 13, -1, 15],
       [16, 17, 18, 19]])

VERWENDEN:

 >>> a[[0,1,3]][:,[0,2]]
array([[ 0,  2],
   [ 4,  6],
   [12, 14]])

ODER:

>>> a[[0,1,3],::2]
array([[ 0,  2],
   [ 4,  6],
   [12, 14]])

Die Verwendung np.ix_ist der bequemste Weg, dies zu tun (wie von anderen beantwortet), aber hier ist ein anderer interessanter Weg, dies zu tun:

>>> rows = [0, 1, 3]
>>> cols = [0, 2]

>>> a[rows].T[cols].T

array([[ 0,  2],
       [ 4,  6],
       [12, 14]])