Comment extraire des éléments d'un tableau NumPy à l'aide de l'indexation de tableaux entiers ?

Extraction indexée avec un tableau NumPy

Lorsque vous travaillez avec des tableaux multidimensionnels dans NumPy, il est souvent nécessaire d'extraire des éléments spécifiques basés sur des indices. À cet égard, certains scénarios peuvent survenir où un tableau (B) sert d'index pour la deuxième dimension d'un autre tableau (A).

Considérons un exemple avec deux tableaux A et B :

A = np.array([[0,1], 
              [2,3], 
              [4,5]])
B = np.array([[1],
              [0],
              [1]], dtype='int')

Notre objectif est d'extraire un élément de chaque ligne de A, où l'index de l'élément est déterminé par l'élément correspondant dans B. Le résultat souhaité est :

C = np.array([[1],
              [2],
              [5]])

Solution utilisant l'indexation de tableaux entiers :

Au lieu de diffuser B, nous pouvons utiliser l'indexation de tableaux purement entiers de NumPy :

A[np.arange(A.shape[0]),B.ravel()]

Ici, np.arange(A.shape[0]) crée un tableau d'indices correspondant au nombre de lignes dans A, tandis que B.ravel() aplatit B en un tableau 1D. En les combinant, nous pouvons extraire les éléments appropriés de A.

Exemple d'exécution :

print(A)
print(B)
print(A[np.arange(A.shape[0]),B.ravel()])

Sortie :

[[0 1]
 [2 3]
 [4 5]]
[[1]
 [0]
 [1]]
[1 2 5]

Cette approche fournit une solution concise et efficace à notre problème d'extraction.

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!