python - 如何在整个 numpy 矩阵中获取最大(顶部)N 个值

Question

我想在整个 numpy 矩阵中获取前 N 个(最大)参数和值，而不是在单个维度(行/列)中。

示例输入(N=3):

import numpy as np
mat = np.matrix([[9,8, 1, 2], [3, 7, 2, 5], [0, 3, 6, 2], [0, 2, 1, 5]])

print(mat)

[[9 8 1 2]
 [3 7 2 5]
 [0 3 6 2]
 [0 2 1 5]]

期望的输出:[9, 8, 7]

由于 max 不能跨单个维度传递，因此按行或按列是行不通的。

# by rows, no 8
np.squeeze(np.asarray(mat.max(1).reshape(-1)))[:3]
array([9, 7, 6])

# by cols, no 7
np.squeeze(np.asarray(mat.max(0)))[:3]
array([9, 8, 6])

我有可用的代码，但对我来说看起来真的很笨重。

# reshape into single vector
mat_as_vector = np.squeeze(np.asarray(mat.reshape(-1)))

# get top 3 arg positions
top3_args = mat_as_vector.argsort()[::-1][:3]

# subset the reshaped matrix
top3_vals = mat_as_vector[top3_args]

print(top3_vals)

array([9, 8, 7])

希望能有任何更短的方式/更有效的方式/神奇的 numpy 函数来做到这一点!

Answer 1

使用 numpy.partition() 为此目的执行完全排序要快得多:

np.partition(np.asarray(mat), mat.size - N, axis=None)[-N:]

假设N<=mat.size .

如果您还需要对最终结果进行排序(除了排在顶部 N 之外)，那么您需要对之前的结果进行排序(但大概您将对比原始数组更小的数组进行排序):

np.sort(np.partition(np.asarray(mat), mat.size - N, axis=None)[-N:]) 

如果您需要将结果从大到小排序，post-pend [::-1]到上一个命令:

np.sort(np.partition(np.asarray(mat), mat.size - N, axis=None)[-N:])[::-1]