Wie kann man mit einem (N-1)-dimensionalen Array effizient auf ein N-dimensionales Array zugreifen?

Greifen Sie mit einem (N-1)-dimensionalen Array auf ein N-dimensionales Array zu

Gegeben sei ein N-dimensionales Array a und ein (N- 1)-dimensionales Array IDX, eine häufige Aufgabe besteht darin, auf Elemente in einem durch die Indizes in IDX angegebenen Element zuzugreifen. Dies kann nützlich sein, um Operationen wie das Finden von Maxima oder das Abrufen bestimmter Werte durchzuführen.

Elegante Lösung mit erweiterter Indizierung

Eine elegante Lösung besteht darin, erweiterte Indizierung mit der Ogrid-Funktion von NumPy zu verwenden :

<code class="python">m, n = a.shape[1:]
I, J = np.ogrid[:m, :n]
a_max_values = a[idx, I, J]
b_max_values = b[idx, I, J]</code>

Dadurch wird ein Netzgitter aus Indizes erstellt und zur Indizierung in a und b verwendet, was zu Arrays mit den entsprechenden Werten führt.

Allgemeiner Fall mit Funktion

Für eine allgemeinere Lösung, die für jede angegebene Achse funktioniert, können wir eine Funktion definieren:

<code class="python">def argmax_to_max(arr, argmax, axis):
    new_shape = list(arr.shape)
    del new_shape[axis]

    grid = np.ogrid[tuple(map(slice, new_shape))]
    grid.insert(axis, argmax)

    return arr[tuple(grid)]</code>

Diese Funktion benötigt ein Array, seinen Argmax entlang einer angegebenen Achse und die Achse selbst . Anschließend wird ein Meshgrid erstellt und zum Extrahieren der entsprechenden Elemente verwendet.

Vereinfachte Indizierung mit benutzerdefinierter Funktion

Um den Indizierungsprozess weiter zu vereinfachen, können wir eine Hilfsfunktion erstellen das ein Raster von Indizes generiert:

<code class="python">def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)</code>

Diese Funktion gibt ein Tupel von Indizes zurück, das direkt zum Indizieren in Eingabearrays verwendet werden kann:

<code class="python">axis = 0
a_max_values = a[all_idx(idx, axis=axis)]
b_max_values = b[all_idx(idx, axis=axis)]</code>

Das obige ist der detaillierte Inhalt vonWie kann man mit einem (N-1)-dimensionalen Array effizient auf ein N-dimensionales Array zugreifen?. Für weitere Informationen folgen Sie bitte anderen verwandten Artikeln auf der PHP chinesischen Website!