如何跨列有效地部分 argsort Pandas 数据框

Posted 2023-03-11

【中文标题】如何跨列有效地部分 argsort Pandas 数据框

【英文标题】：How to efficiently partial argsort Pandas dataframe across columns

【发布时间】：2019-05-08 03:18:33

【问题描述】：

我想根据每行最大的 3 个值用列标签替换值。让我们假设这个输入：

   p1  p2  p3  p4
0   0   9   1   4
1   0   2   3   4
2   1   3  10   7
3   1   5   3   1
4   2   3   7  10

鉴于n = 3，我正在寻找：

  Top1 Top2 Top3
0   p2   p4   p3
1   p4   p3   p2
2   p3   p4   p2
3   p2   p3   p1
4   p4   p3   p2

我不担心重复，例如对于索引3，Top3 可以是'p1' 或 'p4'。

尝试 1

我的第一次尝试是使用np.ndarray.argsort 进行完整排序：

res = pd.DataFrame(df.columns[df.values.argsort(1)]).iloc[:, len(df.index): 0: -1]

但实际上我有超过 4 列，这将是低效的。

尝试 2

接下来我尝试了np.argpartition。但是由于每个分区中的值没有排序，这需要一个后续排序：

n = 3
parts = np.argpartition(-df.values, n, axis=1)[:, :-1]
args = (-df.values[np.arange(df.shape[0])[:, None], parts]).argsort(1)

res = pd.DataFrame(df.columns[parts[np.arange(df.shape[0])[:, None], args]],
                   columns=[f'Topi' for i in range(1, n+1)])

事实上，这比第一次尝试更大的数据帧时慢。有没有更有效的方法可以利用部分排序？您可以使用以下代码进行基准测试。

基准测试

# Python 3.6.0, NumPy 1.11.3, Pandas 0.19.2

import pandas as pd, numpy as np

df = pd.DataFrame('p1': [0, 0, 1, 1, 2],
                   'p2': [9, 2, 3, 5, 3],
                   'p3': [1, 3, 10, 3, 7],
                   'p4': [4, 4, 7, 1, 10])

def full_sort(df):
    return pd.DataFrame(df.columns[df.values.argsort(1)]).iloc[:, len(df.index): 0: -1]

def partial_sort(df):
    n = 3
    parts = np.argpartition(-df.values, n, axis=1)[:, :-1]
    args = (-df.values[np.arange(df.shape[0])[:, None], parts]).argsort(1)
    return pd.DataFrame(df.columns[parts[np.arange(df.shape[0])[:, None], args]])

df = pd.concat([df]*10**5)

%timeit full_sort(df)     # 86.3 ms per loop
%timeit partial_sort(df)  # 158 ms per loop

【问题讨论】：

列数越多，使用 argparition 的运气可能会更好。有 4 列，似乎不值得。

@Divakar，是的，你可能是对的。我只是认为可能有一些聪明的 NumPy 魔法（pre-sorted np.argpartition?）可以在这里提供帮助。也许它不存在。

那么，通常你有多少列？

@Divakar，约 30 列。对于它的价值，即使这个数字full_sort 工作得更快。随着列数的增加，似乎 2 倍的性能差异仍然存在。

我相信这可能有用：***.com/questions/42184499/…

【参考方案1】：

有了相当多的列，我们可以将np.argpartition 与一些slicing 和indexing 一起使用，就像这样-

def topN_perrow_colsindexed(df, N):
    # Extract array data
    a = df.values

    # Get top N indices per row with not necessarily sorted order
    idxtopNpart = np.argpartition(a,-N,axis=1)[:,-1:-N-1:-1]

    # Index into input data with those and use argsort to force sorted order
    sidx = np.take_along_axis(a,idxtopNpart,axis=1).argsort(1)
    idxtopN = np.take_along_axis(idxtopNpart,sidx[:,::-1],axis=1)    

    # Index into column values with those for final output
    c = df.columns.values
    return pd.DataFrame(c[idxtopN], columns=[['Top'+str(i+1) for i in range(N)]])

示例运行 -

In [65]: df
Out[65]: 
   p1  p2  p3  p4
0   0   9   1   4
1   0   2   3   4
2   1   3  10   7
3   1   5   3   1
4   2   3   7  10

In [66]: topN_perrow_colsindexed(df, N=3)
Out[66]: 
  Top1 Top2 Top3
0   p2   p4   p3
1   p4   p3   p2
2   p3   p4   p2
3   p2   p3   p4
4   p4   p3   p2

时间安排 -

In [143]: np.random.seed(0)

In [144]: df = pd.DataFrame(np.random.rand(10000,30))

In [145]: %timeit full_sort(df)
     ...: %timeit partial_sort(df)
     ...: %timeit topN_perrow_colsindexed(df,N=3)
100 loops, best of 3: 7.96 ms per loop
100 loops, best of 3: 13.9 ms per loop
100 loops, best of 3: 5.47 ms per loop

In [146]: df = pd.DataFrame(np.random.rand(10000,100))

In [147]: %timeit full_sort(df)
     ...: %timeit partial_sort(df)
     ...: %timeit topN_perrow_colsindexed(df,N=3)
10 loops, best of 3: 34 ms per loop
10 loops, best of 3: 56.1 ms per loop
100 loops, best of 3: 13.6 ms per loop