如何在Pandas DataFrame中按规则补全缺失季度行并设置status

实现思路 #

要搞定“连续多个未使用季度只在第一个季度加'out'行”这个需求，核心思路是先列出所有文件和季度的完整组合，再筛选出连续缺失段的首个季度，最后将补充行与原数据合并即可。

具体代码实现 #

步骤1：生成所有文件-季度的完整组合 #

先提取原数据中的唯一文件名，和给定季度列表做全量配对，确保每个文件都对应所有季度：

import pandas as pd

# 原测试数据
df_test = pd.DataFrame(data=None, columns=['file', 'quarter', 'status'])
df_test.file = ['file_1', 'file_1', 'file_2', 'file_2', 'file_3']
df_test.quarter = ['2022q4', '2023q2', '2022q3', '2022q4', '2023q1']
df_test.status = ['in', 'in', 'in', 'in', 'in']

# 固定季度列表
quarters = ['2022q3', '2022q4', '2023q1', '2023q2']

# 生成全量的文件-季度组合
all_files = df_test['file'].unique()
full_df = pd.MultiIndex.from_product([all_files, quarters], names=['file', 'quarter']).to_frame(index=False)

步骤2：合并原数据，标记缺失行 #

将全量组合与原数据左连接，找出原数据中没有记录的季度（即需要补充的行）：

merged_df = pd.merge(full_df, df_test, on=['file', 'quarter'], how='left')
# 给缺失行打标记
merged_df['is_missing'] = merged_df['status'].isna()

步骤3：分组筛选连续缺失的首个季度 #

按文件分组处理，先按季度排序，再只保留连续缺失段的第一行：

def filter_first_missing(group):
    # 按季度排序（确保是时间顺序，这里quarters本身已排好序）
    group = group.sort_values('quarter')
    # 标记连续缺失的起始行
    group['is_first_missing'] = False
    # 第一行如果是缺失的，直接标记
    if group.iloc[0]['is_missing']:
        group.iloc[0, group.columns.get_loc('is_first_missing')] = True
    # 从第二行开始，判断当前行缺失且上一行不缺失的，就是连续缺失的起点
    for i in range(1, len(group)):
        if group.iloc[i]['is_missing'] and not group.iloc[i-1]['is_missing']:
            group.iloc[i, group.columns.get_loc('is_first_missing')] = True
    # 保留原有的'in'行，以及标记好的首个缺失行
    return group[(~group['is_missing']) | group['is_first_missing']]

# 分组处理每个文件
result_df = merged_df.groupby('file').apply(filter_first_missing).reset_index(drop=True)

步骤4：填充状态并整理结果 #

将筛选出的缺失行状态填充为'out'，删除辅助列后排序：

# 填充缺失的状态为'out'
result_df['status'] = result_df['status'].fillna('out')
# 删除辅助列
result_df = result_df.drop(['is_missing', 'is_first_missing'], axis=1)
# 按文件和季度排序，匹配期望输出
result_df = result_df.sort_values(['file', 'quarter']).reset_index(drop=True)

print(result_df)

最终输出 #

运行代码后得到的结果与期望一致：

file quarter status
0  file_1  2022q4     in
1  file_1  2023q1    out
2  file_1  2023q2     in
3  file_2  2022q3     in
4  file_2  2022q4     in
5  file_2  2023q1    out
6  file_3  2023q1     in
7  file_3  2023q2    out

性能优化提示 #

如果处理的是大型DataFrame，循环效率较低，推荐用向量化写法替代循环，速度提升明显：

def filter_first_missing_vectorized(group):
    group = group.sort_values('quarter')
    # 用shift()判断：当前行缺失，且前一行不缺失（前一行不存在则用False填充）
    group['is_first_missing'] = group['is_missing'] & ~group['is_missing'].shift(fill_value=False)
    return group[(~group['is_missing']) | group['is_first_missing']]

替换上述filter_first_missing函数即可，逻辑完全一致，处理大数据集时更高效。

内容的提问来源于stack exchange，提问作者Marcus K.