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]])