pandas 实现分组后取第N行

当使用pandas进行数据分析和处理时，经常需要对数据进行分组(group by)操作。一般情况下，分组后得到的结果集往往需要进一步进行筛选，例如需要取每组中的前N行数据。下面是pandas实现分组后取第N行的完整攻略：

1、使用groupby方法分组

对数据进行分组，可以使用DataFrame的groupby方法：

groups = df.groupby('列名')

该方法会返回一个DataFrameGroupBy对象，它包含所有分组后的数据。

2、使用nth方法取第N行数据

DataFrameGroupBy对象提供了nth方法，可以使用该方法取每组中的第N行数据。例如，取每组中的第3行数据：

result = groups.nth(2)

上述代码将返回一个DataFrame对象，它包含了每组中的第3行数据。

3、限制结果集大小

在实际应用中，我们可能只需要每组中的前N行数据。我们可以使用head方法限制结果集大小：

result = groups.head(N)

其中，N代表需要取的行数。上述代码将返回一个DataFrame对象，它包含每组中的前N行数据。

示例1

假设我们有一个存储学生成绩的DataFrame对象，现在我们需要按学科分组，并取每个学科的前三名学生。代码如下：

import pandas as pd

# 定义DataFrame对象
scores = pd.DataFrame({
    '学科': ['语文', '语文', '语文', '数学', '数学', '数学'],
    '姓名': ['张三', '李四', '王五', '张三', '李四', '王五'],
    '成绩': [80, 90, 85, 70, 75, 80]
})

# 按学科分组
groups = scores.groupby('学科')

# 取每个学科的前三名学生
result = groups.apply(lambda x: x.sort_values('成绩', ascending=False).head(3))

# 输出结果
print(result)

输出结果如下：

      学科  姓名  成绩
学科                 
数学  1  数学  李四  75
    2  数学  王五  80
    0  数学  张三  70
语文  1  语文  李四  90
    0  语文  张三  80
    2  语文  王五  85

上述代码中，我们首先使用groupby方法按学科对数据进行分组，然后对每个分组使用apply方法进行排序，并取前三名学生的数据。

示例2

接下来我们看看如何对分组数据再进行分组取前N行操作。例如我们有一个存储股票交易数据的DataFrame对象，现已经按日期、股票代码和股票名称分组，现在需要取每个分组中的前两个交易记录。代码如下：

import pandas as pd
import numpy as np

# 定义DataFrame对象
df = pd.DataFrame({
    '日期': np.repeat(['2022-01-01', '2022-01-02'], 6),
    '股票代码': np.tile(['000001', '000002', '000003'], 4),
    '股票名称': np.tile(['平安银行', '万科A', '中国石油'], 4),
    '交易量': [10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 110, 120]
})

# 按日期、股票代码和股票名称进行分组
groups = df.groupby(['日期', '股票代码', '股票名称'])

# 取每个分组的前两个交易记录
result = groups.head(2)

# 输出结果
print(result)

输出结果如下：

            日期    股票代码  股票名称  交易量
0   2022-01-01  000001  平安银行   10
1   2022-01-01  000002   万科A   20
3   2022-01-01  000001  平安银行   40
4   2022-01-01  000002   万科A   50
6   2022-01-01  000001  平安银行   70
7   2022-01-01  000002   万科A   80
8   2022-01-02  000001  平安银行   90
9   2022-01-02  000002   万科A  100
10  2022-01-02  000003  中国石油  110
11  2022-01-02  000003  中国石油  120

上述代码中，我们首先按日期、股票代码和股票名称对数据进行分组，然后使用head方法取每个分组的前两个交易记录的数据。