Существует ли функция numpy, которая возвращает матрицу индексов массива заданного размера?

Question

У меня есть массив 28x28 numpy, и я хочу создать новый массив numpy кортежей, каждый кортеж представляет индекс в исходном массиве, что-то вроде: [(0, 0), (0, 1) , ..., (0, 28), (1, 0), (1, 1), ..., (1, 28), ... и т. Д.

Есть ли простой способ сделать это?

Answer 1

Существует np.ndindex, который можно использовать для перебора N-мерных индексов по форме массива:

a = np.random.rand(28,28)

list(np.ndindex(a.shape))
# [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0,...

Answer 2

Точно так же, как альтернатива, поскольку вам нужен список кортежей, вы также можете рассмотреть возможность не использовать NumPy:

from itertools import product

a = ...  # Your array
idx_tuples = product(*map(range, a.shape))

Тогда list(idx_tuples) выдаст вам список, который вы хотите. Это, вероятно, будет медленнее, единственное довольно незначительное преимущество заключается в том, что, если вам нужен только этот список для повторения его один раз, вы можете выполнить for idx in idx_tuples: для генерации кортежей во время итерации, сохраняя память для их предварительного создания.

Answer 3

Кроме того np.ndindex() (что является самым простым в данном случае), вы можете использовать np.unravel_index() и np.ravel_multi_index() до go назад и вперед от линейного и многоиндексный.

Например:

import numpy as np


shape = (3, 5)
arr = np.random.rand(*shape)
print([np.unravel_index(i, shape) for i in range(np.prod(shape))])
# [(0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (1, 0), (1, 1), (1, 2), (1, 3), (1, 4), (2, 0), (2, 1), (2, 2), (2, 3), (2, 4)]