python - 在numpy数组的特定行中，如何查找前3个最大值的列索引

Question

我有一个数组 X：

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

我想要索引为 1 的前 3 个最大值的索引。结果将是：

[0,2,1]

我对 Python 比较陌生。我尝试使用 argsort 执行此操作，但无法针对某一特定行执行此操作。

标签： pythonarraysnumpy

Answer 1

您可以使用argsorton axis=1（按行），然后提取每行的最后 3 个索引：

X.argsort(axis=1)[:,:-4:-1]

#[[2 0 1]
# [0 2 1]
# [0 2 1]
# [2 1 3]
# [3 0 1]]