在numpy中获取匹配条件的行数答案

【问题标题】：Get row numbers of rows matching a condition in numpy在numpy中获取匹配条件的行数
【发布时间】：2015-05-09 23:09:41
【问题描述】：

假设我有一个像这样的 numpy 数组：

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

我想检查第二个元素是否 == 2。

我知道我可以做到：

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

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

【问题讨论】：

标签： python arrays numpy

【解决方案1】：

使用np.where 返回索引：

In [79]:

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

【讨论】：

【解决方案2】：

这是我使用的函数，它只提供行号：

import numpy as np

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

row_locator = a[:,1]==2

row_values = []

for x in range(0, len(a)):
    if row_locator[x]:
        row_values.append(x)

print(row_values)

>>>[0, 3]

【讨论】：