获取与numpy中的条件匹配的行的行号

问题描述

假设我有一个类似numpy的数组:

Suppose I have a numpy array like:

a = array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9],
       [3, 2, 1]])

我要检查第二个元素是否==2.

I want to check if the second element == 2.

我知道我可以做到:

>>> a[:,1]==2
array([ True, False, False,  True], dtype=bool)

返回布尔值.我的问题是，如何获得条件为真的行的行号?在此示例中，我想返回array([0, 3])，因为第0行和第3行匹配第二个元素== 2的条件.

returning booleans. My question is, how do I get the row numbers of the rows where the condition is true? In this example I would want to get back array([0, 3]) because the 0th and 3rd rows match the condition second element == 2.

推荐答案

使用 np.where 返回索引:

Use np.where to return the indices:

In [79]:

np.where(a[:,1]==2)
Out[79]:
(array([0, 3], dtype=int64),)

这篇关于获取与numpy中的条件匹配的行的行号的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！