如何对Pandas DataFrame分组切片并按需补全，保留原顺序与索引？

现有如下Pandas DataFrame：

import pandas as pd
df = pd.DataFrame({'mycol': ['A', 'B', 'A', 'B', 'B', 'C', 'A', 'C', 'A', 'A']})

print(df)

输出：

mycol
0     A
1     B
2     A
3     B
4     B
5     C
6     A
7     C
8     A
9     A

其中分组A出现5次、B出现3次、C出现2次，需求如下：

• 每个分组仅保留3条数据：A组超出的2条（索引8、9）删除；C组缺少的1条新增（索引设为10）
• 严格保留原行顺序
• 索引需可追踪位置

本人编写的代码出现索引混乱、多余NaN行的问题：

def func(group):
    df2 = pd.DataFrame(None, index=[max(group.index)+1], columns=['mycol'])
    result = pd.concat([group.iloc[:3], df2])
    result['newcol'] = [group.name + str(i+1) for i in range(len(result))] 
    return result
    
final = df.groupby('mycol').apply(func).droplevel(0)
print(final)

错误输出：

mycol newcol
0      A     A1
2      A     A2
6      A     A3
10   NaN     A4
1      B     B1
3      B     B2
4      B     B3
5    NaN     B4
5      C     C1
7      C     C2
8    NaN     C3

期望输出：

mycol  newcol
0      A      A1
1      B      B1
2      A      A2
3      B      B2
4      B      B3
5      C      C1
6      A      A3
7      C      C2
10   NaN      C3

问题根源 #

原代码的问题在于：

1. 对所有分组强制新增一行NaN（包括不需要补行的B组）
2. groupby.apply按分组拼接结果，打乱了原数据的行顺序

修正代码 #

import pandas as pd

df = pd.DataFrame({'mycol': ['A', 'B', 'A', 'B', 'B', 'C', 'A', 'C', 'A', 'A']})

# 1. 给每个分组内的行标记序号，筛选保留前3条有效行
df['group_seq'] = df.groupby('mycol').cumcount() + 1
valid_rows = df[df['group_seq'] <= 3].copy()

# 2. 统计各分组行数，找出需要补行的分组
group_counts = df['mycol'].value_counts()
need_fill = group_counts[group_counts < 3].index

# 3. 生成补全行，索引从原最大索引（9）开始递增
max_idx = df.index.max()
fill_rows = []
for col in need_fill:
    fill_num = 3 - group_counts[col]
    for _ in range(fill_num):
        max_idx += 1
        fill_rows.append({'mycol': None, 'group_seq': group_counts[col] + 1, 'index': max_idx})

fill_df = pd.DataFrame(fill_rows).set_index('index')

# 4. 合并有效行与补全行，按索引排序后生成newcol
final = pd.concat([valid_rows, fill_df]).sort_index()
final['newcol'] = final['mycol'].fillna(final['mycol'].mode()[0]) + final['group_seq'].astype(str)
final = final.drop('group_seq', axis=1)

print(final)

输出结果：

mycol newcol
0      A     A1
1      B     B1
2      A     A2
3      B     B2
4      B     B3
5      C     C1
6      A     A3
7      C     C2
10   NaN     C3

代码说明 #

1. cumcount()给每个分组内的行按原顺序编号，筛选前3条有效行，保留原索引和行顺序
2. 仅对行数不足3的分组生成补全行，索引延续原数据的最大索引，确保可追踪
3. 合并后按索引排序，严格还原原数据的行顺序
4. 通过fillna结合分组名和序号生成newcol，补全行的newcol自动匹配对应分组

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