Pandas：将分组变量扩展至包含数据框下一组的前n条观测值

针对你的需求——让每个分组的标签覆盖自身所有行，同时覆盖下一个分组的前n条行，最后一个分组的标签仅保留自身未被前序分组覆盖的部分——可以通过以下步骤实现：

代码实现 #

import pandas as pd

# 初始化示例DataFrame
egdf = pd.DataFrame({
    'group': ['A']*6 + ['B']*6 + ['C']*5 + ['D']*6,
    'value': list(range(1,7))*2 + list(range(1,6)) + list(range(1,7))
})

# 定义需要覆盖下一组的行数n
n = 2
# 获取按出现顺序排列的唯一分组列表（pandas unique()保留元素首次出现顺序）
groups = egdf['group'].unique()

# 初始化结果列为原分组列
egdf['output_wanted'] = egdf['group'].copy()

# 遍历每个分组（跳过最后一个），修改下一组前n行的标签
for i in range(len(groups) - 1):
    current_group = groups[i]
    next_group = groups[i+1]
    # 筛选下一组的前n行索引
    next_topn_indices = egdf[egdf['group'] == next_group].head(n).index
    # 将这些行的结果标签替换为当前分组名
    egdf.loc[next_topn_indices, 'output_wanted'] = current_group

验证结果 #

执行后egdf['output_wanted']的结果与示例完全一致：

print(egdf['output_wanted'].tolist())
# 输出：['A']*8 + ['B']*6 + ['C']*5 + ['D']*4

逻辑说明 #

1. 分组顺序保留：通过unique()获取原数据中分组的出现顺序，确保后续处理的分组顺序与原数据一致。
2. 初始标签赋值：先将结果列设为原分组列，后续仅修改需要覆盖的行。
3. 覆盖下一组前n行：循环处理每个分组（除最后一个），找到下一组的前n行索引，将这些行的标签替换为当前分组名，实现分组扩展覆盖的需求。
4. 特殊情况处理：最后一个分组没有后续分组，因此其标签自动保留自身未被前序分组覆盖的部分（即从第n+1行开始的剩余行）。

注意事项 #

如果原数据中同分组的行不是连续排列的，需要先按group列排序并重置索引，确保同分组行连续，否则head(n)无法正确取到下一组的前n行：

# 若分组不连续，先排序重置索引
egdf = egdf.sort_values('group').reset_index(drop=True)

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