获取pandas数据框中每一列的前k个元素的索引的快速方法

Posted 2023-03-12

【中文标题】获取pandas数据框中每一列的前k个元素的索引的快速方法

【英文标题】：fast way to get index of top-k elements of every column in a pandas dataframe

【发布时间】：2015-11-18 06:20:30

【问题描述】：

我有一个非常大的 pandas 数据框，大约有 500,000 列。每列大约有 500 个元素长。对于每一列，我需要检索列中前 k 个元素的（索引、列）位置。

所以，如果 k 等于 2，这是我的数据框：

  A  B  C  D
w 4  8  10 2
x 5  1  1  6 
y 9  22 25 7 
z 15 5  7  2

我想回来：

[(A,y),(A,z),(B,w),(B,y),(C,w),(C,y),(D,x),(D,y)]

请记住，我有大约 500,000 列，所以速度是我最关心的问题。有没有一种合理的方法可以在我的机器上花费一整周的时间？最快的方法是什么——即使它对于我拥有的数据量来说足够快？

感谢您的帮助！

【问题讨论】：

这里的问题是你真正想要返回什么，因为构造一个元组列表会很昂贵，例如df.apply(lambda x: x.sort(inplace=False, ascending=False)[:2]) 会返回你想要的，但是NaN 的值有点笨拙，可能不是你想要的

【参考方案1】：

Pandas 有一个高效的nlargest 操作，您可以使用它比完整排序更快。跨 500,000 列应用仍需要一段时间。

In [1]: df = pd.DataFrame(data=np.random.randint(0, 100, (200, 500000)), 
                          columns=range(500000), index=range(200))

In [2]: %time np.array([df[c].nlargest(2).index.values for c in df])
Wall time: 2min 57s
Out[2]: 
array([[171,   1],
       [ 42,  78],

正如@EdChum 所说，您可能不想存储为元组，使用两个数组或其他策略会更有效率。

【讨论】：

这应该比临时排序 +1 更快，这里的问题是apply 试图返回一个与原始 df 形状相同的 df，除非您获取原始值并返回一些其他数据结构，例如您的答案

【参考方案2】：

我认为numpy 对此有一个很好的解决方案，速度很快，您可以根据需要格式化输出。

In [2]: df = pd.DataFrame(data=np.random.randint(0, 1000, (200, 500000)), 
                      columns=range(500000), index=range(200))

In [3]: def top_k(x,k):
             ind=np.argpartition(x,-1*k)[-1*k:]
             return ind[np.argsort(x[ind])]

In [69]: %time np.apply_along_axis(lambda x: top_k(x,2),0,df.as_matrix())
CPU times: user 5.91 s, sys: 40.7 ms, total: 5.95 s
Wall time: 6 s

Out[69]:
array([[ 14,  54],
       [178, 141],
       [ 49, 111],
       ...,
       [ 24, 122],
       [ 55,  89],
       [  9, 175]])

与 pandas 解决方案相比相当快（IMO 更干净，但我们在这里追求速度）：

In [41]: %time np.array([df[c].nlargest(2).index.values for c in df])
CPU times: user 3min 43s, sys: 6.58 s, total: 3min 49s
Wall time: 4min 8s

Out[41]:
array([[ 54,  14],
       [141, 178],
       [111,  49],
       ...,
       [122,  24],
       [ 89,  55],
       [175,   9]])

这些列表彼此顺序相反（您可以通过在 numpy 版本中反转排序来轻松解决此问题）

请注意，在示例中，由于随机 int 生成，我们可能有多个 k 相等且最大的值，因此返回的索引可能在所有方法中不一致，但所有方法都会产生有效结果（您将得到 @987654326 @与列中最大值匹配的索引）

【讨论】：

这是一种非常快速的方法，我使用了它。不幸的是，它不提供 pandas 索引，只提供索引号。但是很容易修改代码来解决这个问题。谢谢！

干得好，如果您认为它会帮助其他人，请随时编辑答案以包含到熊猫索引的映射。