numpy.ndindex

class numpy.ndindex(*shape)

Класс ndindex() создает объект итератора, который возвращающает индексы элементов в массиве заданной формы.

Параметры:
*shape - целые положительные числа.
Каждое число определяет размер соответствующей оси массива. Обратите внимание, что это именно числа, перечисленные через запятую, а непоследовательность чисел.
Смотрите так же:
ndenumerate, flatiter, flat, flatten


Примеры

>>> import numpy as np
>>> 
>>> for ind in np.ndindex(2, 2, 3):
...     print(ind)
... 
(0, 0, 0)
(0, 0, 1)
(0, 0, 2)
(0, 1, 0)
(0, 1, 1)
(0, 1, 2)
(1, 0, 0)
(1, 0, 1)
(1, 0, 2)
(1, 1, 0)
(1, 1, 1)
(1, 1, 2)

Возвращаемые индексы могут быть использованы для индексирования любого массива соответствующей формы:

>>> a = np.arange(12).reshape(2, 2, 3)
>>> 
>>> for ind in np.ndindex(2, 2, 3):
...     print(ind, a[ind])
... 
(0, 0, 0) 0
(0, 0, 1) 1
(0, 0, 2) 2
(0, 1, 0) 3
(0, 1, 1) 4
(0, 1, 2) 5
(1, 0, 0) 6
(1, 0, 1) 7
(1, 0, 2) 8
(1, 1, 0) 9
(1, 1, 1) 10
(1, 1, 2) 11

Метод next() возвращает следующую итерацию:

>>> ind = np.ndindex(2, 2, 3)
>>> 
>>> ind.next()
(0, 0, 0)
>>> ind.next()
(0, 0, 1)
>>> ind.next()
(0, 0, 2)