Python Numpy 2-мерный массив итераций

Question

Хотел бы построить список индексов в двумерном массиве bool_, где True.

import numpy
arr = numpy.zeros((6,6), numpy.bool_)
arr[2,3] = True
arr[5,1] = True
results1 = [[(x,y) for (y,cell) in enumerate(arr[x].flat) if cell] for x in xrange(6)]
results2 = [(x,y) for (y,cell) in enumerate(arr[x].flat) if cell for x in xrange(6)]

результаты 1:

[[], [], [(2, 3)], [], [], [(5, 1)]]

результаты 2 полностью неверны

Цель:

[(2, 3),(5, 1)]

Есть ли способ сделать это без сглаживания списка впоследствии, или какой-нибудь лучший способ сделать это в целом?

Answer 1

Я думаю, что функция, которую вы ищете, это numpy.where . Вот пример:

>>> import numpy
>>> arr = numpy.zeros((6,6), numpy.bool_)
>>> arr[2,3] = True
>>> arr[5,1] = True
>>> numpy.where(arr)
(array([2, 5]), array([3, 1]))

Вы можете превратить это обратно в индекс, подобный этому:

>>> numpy.array(numpy.where(arr)).T
array([[2, 3],
       [5, 1]])

Answer 2

>>> import numpy as np
>>> arr = np.zeros((6,6), np.bool_)
>>> arr[2,3] = True
>>> arr[5,1] = True
>>> np.argwhere(arr)
array([[2, 3],
       [5, 1]])