快速获取pandas数据帧中每列的前k个元素的索引

时间:2015-08-24 18:08:36

标签: sorting pandas indexing columnsorting

我有一个非常大的pandas数据框,大约有500,000列。每列长约500个元素。对于每列,我需要检索列中top-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列,所以速度是我最关心的问题。有没有一种合理的方法可以在我的机器上花费整整一周的时间?什么是最快的方式 - 即使它对我的数据量足够快?

感谢您的帮助!

2 个答案:

答案 0 :(得分:5)

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 :(得分:4)

我认为@Path("foo") public class Foo { @Inject public MyContext message; @GET public String index() throws UnknownHostException { return String.format("%s number of instances: %s", message, MyContext.instances); } 有一个很好的解决方案,速度快,您可以根据需要格式化输出。

numpy

与熊猫解决方案(这是更清洁的IMO,但我们在这里寻求速度)相比非常快:

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

列表的顺序相反(您可以通过在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]]) 版本中反转排序轻松解决此问题)

请注意,在示例中,由于随机int生成,我们可能会有超过numpy个等于和最大值的值,因此返回的索引可能在所有方法中都不一致,但所有方法都会产生有效结果(您将获得{(1}}个与列中的最大值匹配的索引