Obtenir les indices et les valeurs d'un ndarray en NumPy

Question

J'ai un ndarray A de nombre arbitraire de dimensions N. Je veux créer un tableau B de tuples (tableau, ou listes) où les premiers N éléments dans chaque tuple sont l'index et le dernier élément est la valeur de cet index dans A.

Par exemple:

A = array([[1, 2, 3], [4, 5, 6]]) 

Puis

B = [(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)] 

Quelle est la meilleure/la plus rapide de le faire dans NumPy sans for boucles?

Answer 1

Si vous avez Python 3 est une manière très simple (et modérément rapide) serait (en utilisant np.ndenumerate):

>>> import numpy as np 
>>> A = np.array([[1, 2, 3], [4, 5, 6]]) 
>>> [(*idx, val) for idx, val in np.ndenumerate(A)] 
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)] 

Il serait un peu différent si vous voulez que cela fonctionne pour les deux Python 3 et Python 2, car Python 2 n'autorise pas le déballage itérable à l'intérieur d'un tuple littéral. Mais vous pouvez utiliser concaténation tuple (ajout):

>>> [idx + (val,) for idx, val in np.ndenumerate(A)] 
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)] 

---

Si vous voulez rester complètement dans NumPy il serait préférable de créer les indices avec np.mgrid:

>>> grid = np.mgrid[:A.shape[0], :A.shape[1]] # indices! 
>>> np.stack([grid[0], grid[1], A]).reshape(3, -1).T 
array([[0, 0, 1], 
     [0, 1, 2], 
     [0, 2, 3], 
     [1, 0, 4], 
     [1, 1, 5], 
     [1, 2, 6]]) 

Cependant cela nécessiterait une boucle pour le convertir en une liste de tuples ... Mais il serait facile de le convertir en une liste de liste:

>>> np.stack([grid[0], grid[1], A]).reshape(3, -1).T.tolist() 
[[0, 0, 1], [0, 1, 2], [0, 2, 3], [1, 0, 4], [1, 1, 5], [1, 2, 6]] 

La liste des tuples est également possible sans visiblefor -loop:

>>> list(map(tuple, np.stack([grid[0], grid[1], A]).reshape(3, -1).T.tolist())) 
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)] 

Même si il n'y a pas visible for -loop le tolist, list, tuple et map ne cachent un for -loop dans la Couche Python.

---

Pour les tableaux de dimensions vous arbitraires besoin de changer cette dernière approche un peu:

coords = tuple(map(slice, A.shape)) 
grid = np.mgrid[coords] 

# array version 
np.stack(list(grid) + [A]).reshape(A.ndim+1, -1).T 
# list of list version 
np.stack(list(grid) + [A]).reshape(A.ndim+1, -1).T.tolist() 
# list of tuple version 
list(map(tuple, np.stack(list(grid) + [A]).reshape(A.ndim+1, -1).T.tolist())) 

L'approche ndenumerate fonctionnerait pour les tableaux de toutes dimensions sans changement et selon mes timings seulement 2-3 fois plus lent.

Answer 2

Vous pouvez également le faire avec np.ndindex, bien que l'approche de @ Mseifert soit plutôt imbattable en termes de timing et de brièveté. La seule boucle ici est de compresser le générateur de coordonnées avec les valeurs réelles. (Comme dans l'autre réponse.)

def tuple_index(a): 
    indices = np.ndindex(*a.shape) 
    return [(*i, j) for i, j in zip(indices, a.flatten())] 

print(tuple_index(a)) 
[(0, 0, 1), (0, 1, 2), (0, 2, 3), (1, 0, 4), (1, 1, 5), (1, 2, 6)]