Comment indexer efficacement des tableaux à N dimensions avec des tableaux d'index de dimension inférieure ?

Indexation d'un tableau à N dimensions avec un tableau à (N-1) dimensions

Accès à un tableau à N dimensions avec un (N -1)-Le tableau dimensionnel présente un défi lors de la recherche de valeurs alignées sur une dimension spécifique. Les approches conventionnelles utilisant np.argmax peuvent ne pas suffire.

Approche d'indexation avancée

Une indexation élégante peut être obtenue grâce à une indexation avancée à l'aide de np.ogrid. Pour un tableau 3D a et son argmax le long de la première dimension, idx :

import numpy as np

a = np.random.random_sample((3, 4, 4))
idx = np.argmax(a, axis=0)

m, n = a.shape[1:]
I, J = np.ogrid[:m, :n]
a_max_values = a[idx, I, J]

Cette approche crée une grille qui étend efficacement le tableau d'index aux dimensions complètes du tableau d'origine.

Généralisation pour les dimensions arbitraires

Pour une solution plus généralisée, la fonction argmax_to_max() peut être définie :

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)]

Cette fonction prend le tableau d'origine, son argmax, et l'axe souhaité et renvoie les valeurs maximales correspondantes.

Approche alternative pour l'indexation générale

Pour indexer tout tableau à N dimensions avec une dimension (N-1) tableau, la fonction all_idx() est une solution plus simplifiée :

def all_idx(idx, axis):
    grid = np.ogrid[tuple(map(slice, idx.shape))]
    grid.insert(axis, idx)
    return tuple(grid)

En utilisant cette fonction, l'indexation dans le tableau a avec idx le long de l'axe peut être réalisée avec :

axis = 0
a_max_values = a[all_idx(idx, axis=axis)]

Ce qui précède est le contenu détaillé de. pour plus d'informations, suivez d'autres articles connexes sur le site Web de PHP en chinois!