Python：在满足条件时查找二维数组中值的行索引答案

【问题标题】：Python: Finding the row index of a value in 2D array when a condition is metPython：在满足条件时查找二维数组中值的行索引
【发布时间】：2020-12-21 13:11:42
【问题描述】：

我有一个维度为 8500 x 5 的二维数组 PointAndTangent。数据是逐行的，有 8500 个数据行和每行 5 个数据值。满足此条件时，我需要提取第 4 列元素的行索引，对于任何s：

abs(PointAndTangent[:,3] - s) <= 0.005

我只需要上述条件的第一个匹配项的行索引。我尝试使用以下内容：

index = np.all([[abs(s - PointAndTangent[:, 3])<= 0.005], [abs(s - PointAndTangent[:, 3]) <= 0.005]], axis=0)

i = int(np.where(np.squeeze(index))[0])

这是行不通的。我收到以下错误：

i = int(np.where(np.squeeze(index))[0])

TypeError: only size-1 arrays can be converted to Python scalars

我对 Python 中的 NumPy 不是很精通。任何建议都会很棒。我试图避免使用for 循环，因为这是我正在尝试的大型模拟的一小部分。

谢谢！

可能的解决方案 我使用了以下

idx = (np.abs(PointAndTangent[:,3] - s)).argmin()

它似乎工作。它返回第 4 列中与 s 最接近的值的行索引。

【问题讨论】：

你想要np.argwhere(abs(PointAndTangent[:, 4] - s) <= 0.005)吗？
我认为这行不通。对于任何s，我都得到 0。
@ShubhamSubhnil。那是因为第一行匹配
我不太确定是什么绊倒了你，但希望我的回答能有所帮助

标签： arrays numpy numpy-ndarray

【解决方案1】：

你快到了。 np.where 是 numpy 中被滥用最多的函数之一。一半时间，你真的想要np.nonzero，另一半时间，你想直接使用布尔掩码。在您的情况下，您需要np.flatnonzero 或np.argmax：

mask = abs(PointAndTangent[:,3] - s) <= 0.005

mask 是一个一维数组，其中满足条件的地方为一，其他地方为零。您可以使用flatnonzero 获取所有索引并选择第一个：

index = np.flatnonzero(mask)[0]

或者，您可以直接使用argmax 选择第一个：

index = np.argmax(mask)

当没有行满足您的条件时，解决方案的行为会有所不同。三个前者做索引，所以会引发错误。后者会返回零，也可以是真实的结果。

通过将mask 替换为分配给它的表达式，两者都可以写成单行。

【讨论】：