我正在使用2D NumPy数组。我想得到(列,行)索引,或者(x,y)坐标,如果您更喜欢这样,我的2D数组符合布尔条件。
我能解释我想要做的最好的方法是通过一个简单的例子:
>>> a = np.arange(9).reshape(3, 3)
>>> b = a > 4
>>> b
>>> array([[False, False, False],
[False, False, True],
[ True, True, True]], dtype=bool)
此时我现在有一个布尔数组,指示a > 4
。
此时我的目标是获取值为True
的布尔数组的索引。例如,索引(1, 2)
,(2, 0)
,(2, 1)
和(2, 2)
都具有值True。
我的最终目标是最终得到一个索引列表:
>>> indexes = [(1, 2), (2, 0), (2, 1), (2, 2)]
同样,我强调指出上面的代码是一个简单的例子,但是我正在尝试做的事情的应用可能有a > 4
的任意索引,而不是基于arange
和reshape
。
答案 0 :(得分:16)
numpy.where
使用numpy.column_stack
:
>>> np.column_stack(np.where(b))
array([[1, 2],
[2, 0],
[2, 1],
[2, 2]])
答案 1 :(得分:5)
@Ashwini Chaudhary答案的替代方案是numpy.nonzero
>>> a = np.arange(9).reshape(3,3)
>>> b = a > 4
>>> np.nonzero(b)
(array([1, 2, 2, 2]), array([2, 0, 1, 2]))
>>> np.transpose(np.nonzero(b))
array([[1, 2],
[2, 0],
[2, 1],
[2, 2]])
编辑:什么是更快。 nonzero
和where
基本上是等效的,但transpose
在这里证明是错误的(即使它是mentioned in the docs):
In [15]: N = 5000
In [16]: a = np.random.random((N, N))
In [17]: %timeit np.nonzero(a > 0.5)
1 loops, best of 3: 470 ms per loop
In [18]: %timeit np.transpose(np.nonzero(a > 0.5)) # ooops
1 loops, best of 3: 2.56 s per loop
In [19]: %timeit np.where(a > 0.5)
1 loops, best of 3: 467 ms per loop
In [20]: %timeit np.column_stack(np.where(a > 0.5))
1 loops, best of 3: 653 ms per loop