python - 如何根据另一列中的条件查找 3D 数组中的最大列？

Question

我有一个包含 3 列的数组，需要找到第三列小于某个预定义值的第一列的最大值。

输入，例如：

a = [[ 1 2 1], [ 2 3 3], [ 7 10 5], [ 5 1 2], [ 2 3 4]]

期望的输出

最大值（a[:,0]，其中 a[:,2] < 5) = [5 1 2]

这个怎么做？最好尽可能快。

标签： pythonpython-3.xnumpymax

Answer 1

您可以使用np.argmax：

subset = a[a[:, 2] < 5]
result = subset[subset[:, 0].argmax()]

print(result)

输出：

[5 1 2]

这将根据给定条件获取原始数组的子集，使用切片和 找到第一列中具有最高值的行的索引argmax，并使用该值来索引子集。