Python Numpy 2維數組迭代

Question

想將索引列表構建到2維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]])