Функция ndindex

Функция ndindex создает итератор, который генерирует кортежи индексов для массива указанной формы. Это позволяет перебирать все возможные позиции элементов в многомерном массиве. Первым параметром функция принимает форму массива в виде кортежа или целых чисел.

Синтаксис

np.ndindex(*shape)

Пример

Давайте создадим одномерный массив и переберем все его индексы с помощью ndindex:

import numpy as np

arr = np.array([10, 20, 30, 40])
for idx in np.ndindex(arr.shape):
	print(f"Index {idx}: value = {arr[idx]}")

Результат выполнения кода:

Index (0,): value = 10
Index (1,): value = 20
Index (2,): value = 30
Index (3,): value = 40

Пример

Переберем индексы двумерного массива размером 2 на 3:

import numpy as np

arr = np.array([[1, 2, 3], [4, 5, 6]])
for idx in np.ndindex(2, 3):
	print(f"Index {idx}: value = {arr[idx]}")

Результат выполнения кода:

Index (0, 0): value = 1
Index (0, 1): value = 2
Index (0, 2): value = 3
Index (1, 0): value = 4
Index (1, 1): value = 5
Index (1, 2): value = 6

Пример

С помощью ndindex можно изменять элементы массива по их индексам:

import numpy as np

arr = np.zeros((2, 3), dtype=int)
for idx in np.ndindex(arr.shape):
	arr[idx] = idx[0] + idx[1]
print(arr)

Результат выполнения кода:

[[0 1 2]
 [1 2 3]]

Пример

Функция ndindex также работает с трехмерными массивами:

import numpy as np

arr = np.arange(8).reshape(2, 2, 2)
for idx in np.ndindex(2, 2, 2):
	print(f"arr{idx} = {arr[idx]}")

Результат выполнения кода:

arr(0, 0, 0) = 0
arr(0, 0, 1) = 1
arr(0, 1, 0) = 2
arr(0, 1, 1) = 3
arr(1, 0, 0) = 4
arr(1, 0, 1) = 5
arr(1, 1, 0) = 6
arr(1, 1, 1) = 7

Смотрите также

функцию indices,
которая создает массив индексов для заданной формы
функцию unravel_index,
которая преобразует плоский индекс в многомерный
функцию ix_,
которая создает сетку индексов для векторной индексации
функцию where,
которая возвращает индексы элементов, удовлетворяющих условию