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

Question

我有一个很大的 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列，所以速度是我的首要考虑。有没有一种合理的方法可以在我的计算机上花费整整一周的时间？最快的方法是什么-即使对于我拥有的数据量来说足够快？

谢谢您的帮助!

Answer 1

我认为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且等于max的值，因此返回的索引可能在所有方法中均不一致，但是所有索引都会产生有效的结果(您将获得与最大值匹配的k索引在列中)