Python перечисляет через 2D массив в numpy

Question

Я хочу функцию, которая ведет себя как enumerate, но на массивах numpy.

>>> list(enumerate("hello")) 
[(0, "h"), (1, "e"), (2, "l"), (3, "l"), (4, "o")] 

>>> for x, y, element in enumerate2(numpy.array([[i for i in "egg"] for j in range(3)])): 
     print(x, y, element) 

0 0 e 
1 0 g 
2 0 g 
0 1 e 
1 1 g 
2 1 g 
0 2 e 
1 2 g 
2 2 g 

В настоящее время я использую эту функцию:

def enumerate2(np_array): 
    for y, row in enumerate(np_array): 
     for x, element in enumerate(row): 
      yield (x, y, element) 

Есть ли лучший способ сделать это? Например. встроенная функция (я не мог найти ее) или другое определение, которое в какой-то мере ускоряется.

Answer 1

Вы хотите np.ndenumerate:

>>> for (x, y), element in np.ndenumerate(np.array([[i for i in "egg"] for j in range(3)])): 
...  print(x, y, element) 
... 
(0L, 0L, 'e') 
(0L, 1L, 'g') 
(0L, 2L, 'g') 
(1L, 0L, 'e') 
(1L, 1L, 'g') 
(1L, 2L, 'g') 
(2L, 0L, 'e') 
(2L, 1L, 'g') 
(2L, 2L, 'g')