如何按行和对numpy数组排序并提取前N行

问题核心与错误原因 #

你之前的代码错误点在于axis参数使用错误：指定axis=1/axis=-1是对每行内部的元素做排序，而我们的需求是先计算每行的元素总和，再根据行和的大小对整行做排序，取行和最大的前top_n行。

正确实现代码 #

import numpy as np

def select_rows(arr, top_n):
    """
    该函数选取元素和最大的前top_n行
    """
    # 计算每行的元素总和
    row_sums = np.sum(arr, axis=1)
    # 对行和做降序排序，取前top_n的行索引
    top_row_indices = np.argsort(-row_sums)[:top_n]
    # 按索引从原数组取对应行
    return arr[top_row_indices]

效果验证 #

用你给出的示例矩阵测试：

# 示例输入矩阵
arr = np.array([[ 6,  7,  8,  9, 10, 11],
       [12, 13, 14, 15, 16, 17],
       [18, 19, 20, 21, 22, 23],
       [ 0,  1,  2,  3,  4,  5],
       [24, 25, 26, 27, 28, 29]])

print(select_rows(arr, 3))

输出结果和预期一致：

[[24 25 26 27 28 29]
 [18 19 20 21 22 23]
 [12 13 14 15 16 17]]

内容的提问来源于stack exchange，提问作者Hefe