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

【问题标题】：How to sort numpy array by row sum and extract top N rows如何按行和对numpy数组进行排序并提取前N行
【发布时间】：2021-10-30 12:54:05
【问题描述】：

例如，给定矩阵

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]])

并且 top_n=3，它应该返回

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

在给定输入二维矩阵 arr 的情况下，此函数应返回形状为 (top_n, arr.shape[-1]) 的 np.ndarray。

这是我尝试过的：

def select_rows(arr, top_n):
    """
    This function selects the top_n rows that have the largest sum of entries
    """
    sel_rows = np.argsort(-arr,axis=1)[:top_n]
    
    return sel_rows

我也试过了：

sel_rows = (-arr).argsort(axis=-1)[:, :top_n]

无济于事。

【问题讨论】：

使用- 将数组转换为负数比在最后对数据进行切片效率低。对于小样本，这不是问题，但是在大数组中将所有值转换为负数会稍微慢一些，这通过%%timeit 测试进行了验证。

标签： python arrays numpy indexing

【解决方案1】：

您的代码几乎可以工作，但您需要在排序之前计算每一行的总和。你可以试试这个：

import numpy as np


top_n = 3
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]])

arr[np.argsort(-arr.sum(axis=1))[:top_n]]

它给出：

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

【讨论】：

答案应该说明-的目的是颠倒顺序

【解决方案2】：

你可以使用这个简单的 1-liner a[np.argsort(a.sum(axis=1))[:-top_n-1:-1]]

a.sum(axis=1) 沿轴 1 求和

np.argsort(..., axis=0) argsorts 沿轴 0（axis=0 无论如何都是默认选项，因此可以省略）

...[:-top_n-1:-1] 以相反的顺序选择最后一个 top_n 索引

a[...] 然后抓取行

`%%timeit`比较

# data sample
a = np.random.randint(0, 101, (100000, 1000))

%%timeit
a[np.argsort(a.sum(axis=1))[:-3-1:-1]]
[out]:
9.73 ms ± 122 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
a[np.argsort(-a.sum(axis=1))[:3]]
[out]:
9.9 ms ± 303 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
sorted(a, key=lambda x: sum(x))[:-3-1:-1]
[out]:
1.04 s ± 36.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

【讨论】：

【解决方案3】：

没有numpy，你可以使用内置函数sorted结合参数key：

sorted(A, key=lambda x: sum(x))[:-top_n-1:-1]

【讨论】：

此实现效率极低，不应与 numpy 数组一起使用。对于数组np.random.randint(0, 101, (100000, 100))，这比 numpy 实现慢 107 倍。

%%timeit比较

`%%timeit`比较