python - 如何在python中获取前N个值

Question

我有一个值(value) list

说

df = DataFrame({'key1' : ['a', 'a', 'b', 'b', 'a'],
   ....:                 'key2' : ['one', 'two', 'one', 'two', 'one'],
   ....:                 'data1' : abs(np.random.randn(5)*100),
   ....:                 'data2' : np.random.randn(5)})

所以如果这是我的数据，

我只想返回 data1 的前 3 个值并返回所有 4 列

除了我想到的很多 if 语句之外，最好的方法是什么。

我正在研究 nlargest ，但不确定我该怎么做

========================更新====================== ===

所以如果上面运行会得到这个结果

我想获得只有 rowindex 为 1,2,3 的返回 df，因为它们具有最高的 data1 前 3 排名(98,94,95)

Answer 1

In [271]: df
Out[271]: 
      data1     data2 key1 key2
0 -1.318436  0.829593    a  one
1  0.172596 -0.541057    a  two
2 -2.071856 -0.181943    b  one
3  0.183276 -1.889666    b  two
4  0.558144 -1.016027    a  one

In [272]: df.ix[df['data1'].argsort()[-3:]]
Out[272]: 
      data1     data2 key1 key2
1  0.172596 -0.541057    a  two
3  0.183276 -1.889666    b  two
4  0.558144 -1.016027    a  one

尽管 heapq.nlargest may be theoretically more efficient ，在实践中，即使对于相当大的 DataFrame，argsort 也往往更快:

import heapq
import pandas as pd
df = pd.DataFrame({'key1' : ['a', 'a', 'b', 'b', 'a']*10000,
                 'key2' : ['one', 'two', 'one', 'two', 'one']*10000,
                 'data1' : np.random.randn(50000),
                 'data2' : np.random.randn(50000)})

In [274]: %timeit df.ix[df['data1'].argsort()[-3:]]
100 loops, best of 3: 5.62 ms per loop

In [275]: %timeit df.iloc[heapq.nlargest(3, df.index, key=lambda x: df['data1'].iloc[x])]
1 loops, best of 3: 1.03 s per loop