Pandas：选择每组的前几行

Posted 2023-03-11

【中文标题】Pandas：选择每组的前几行

【英文标题】：Pandas: select the first couple of rows in each group

【发布时间】：2014-01-08 17:02:39

【问题描述】：

我无法解决这个简单的问题，我在这里寻求帮助... 我有如下DataFrame，我想选择每组'a'中的前两行

df = pd.DataFrame('a':pd.Series(['NewYork','NewYork','NewYork','Washington','Washington','Texas','Texas','Texas','Texas']), 'b': np.arange(9))

df
Out[152]: 
            a  b
0     NewYork  0
1     NewYork  1
2     NewYork  2
3  Washington  3
4  Washington  4
5       Texas  5
6       Texas  6
7       Texas  7
8       Texas  8

也就是说，我想要一个如下的输出：

            a  b
0     NewYork  0
1     NewYork  1
2  Washington  3
3  Washington  4
4       Texas  5
5       Texas  6

非常感谢您的帮助。

【参考方案1】：

在 pandas 0.13rc 中，您可以直接使用 head 执行此操作（即无需 reset_index）：

In [11]: df.groupby('id', as_index=False).head(2)
Out[11]: 
    id   value
0    1   first
1    1  second
3    2   first
4    2  second
5    3   first
6    3   third
9    4  second
10   4   fifth
11   5   first
12   6   first
13   6  second
15   7  fourth
16   7   fifth

[13 rows x 2 columns]

注意：正确的索引，即使有这个小例子，这也比以前快得多（有或没有 reset_index）：

# 0.13rc
In [21]: %timeit df.groupby('id', as_index=False).head(2)
1000 loops, best of 3: 279 µs per loop

# 0.12
In [21]: %timeit df.groupby('id', as_index=False).head(2)  # this didn't work correctly
1000 loops, best of 3: 1.76 ms per loop

In [22]: %timeit df.groupby('id').head(2).reset_index(drop=True)
1000 loops, best of 3: 1.82 ms per loop

【参考方案2】：

抱歉，之前好像有人问过类似的问题... Pandas dataframe get first row of each group 现在明白了...

df.groupby('a').head(2).reset_index(drop=True)
Out[165]: 
            a  b
0     NewYork  0
1     NewYork  1
2       Texas  5
3       Texas  6
4  Washington  3
5  Washington  4