Pandas按分组中另一列值筛选行 双列聚合场景下动态取数实现方法

核心实现代码 #

import pandas as pd

# 构造测试数据集（可替换为自己的数据源读取代码）
data = [
    ['Peter', 0, 12, 'A'],
    ['Peter', 0, 13, 'A'],
    ['Peter', 0, 14, 'B'],
    ['Sam', 1, 12, 'A'],
    ['Sam', 1, 15, 'A'],
    ['Sam', 1, 16, 'B'],
    ['Sam', 1, 17, 'B'],
    ['Mary', 2, 20, 'A'],
    ['Mary', 2, 21, 'A'],
    ['Mary', 2, 12, 'A'],
]
df = pd.DataFrame(data, columns=['Name', 'Shift', 'Data', 'Type'])

# 核心处理逻辑
df['group_inner_idx'] = df.groupby(['Name', 'Type']).cumcount()
result = df[df['group_inner_idx'] == df['Shift']].drop(columns='group_inner_idx').reset_index(drop=True)

# 输出结果
print(result)

逻辑说明 #

• groupby(['Name', 'Type']).cumcount()会对每个Name+Type的分组生成从0开始的连续组内行号，刚好匹配需求中Shift为0取第1条、Shift为1取第2条的规则
• 直接筛选组内行号等于Shift值的记录即可，无需动态传参调用nth方法，实现逻辑简单易维护
  运行上述代码后输出的result即为目标结果。

内容的提问来源于stack exchange，提问作者GChan