用同公司前后同期均值填充DataFrame数值列NaN值的优化方案求助

需求说明 #

我有一个包含金融数据的大型DataFrame，包含分类列与标准数值列。需要实现：

• 筛选出NaN数量少于20000的数值列
• 对这些列中的NaN，用同一家公司（cik）对应周期的前一期和后一期数值的均值填充（例如2022-Q3的NaN用2022-Q2和2022-Q4的均值填充）

原代码问题 #

原代码通过嵌套循环遍历每行数据，每次查询前/后值时都要执行过滤、排序操作，时间复杂度极高，导致运行耗时超20小时，且逻辑难以验证：

def average_col(merged_df_nan):
    #Filling NaN values for columns with less than 20k NaNs with the average of the period before and after
    na_counts = merged_df_nan.isna().sum()
    columns_to_fill2 = na_counts[(na_counts < 20000) & (~na_counts.index.isin(['Text', 'gvkey', 'datadate', 'prccq', 'naics', 'sic', 'processed_text', 'stemmed_text', 'sentiment_score', 'fyr', 'indfmt','consol','popsrc','tic','conm','apdedateq','fdateq','rdq']))].index
    
    #Iterate over the filtered columns and fill NA values with the average of previous and next non-NA values
    for column in tqdm(columns_to_fill2, desc = 'Columns done'):
        na_mask = merged_df_nan[column].isna()
        for index, row in merged_df_nan.iterrows():
            if na_mask[index]:
                fyearq = row['fyearq']
                fqtr = row['fqtr']
                cik = row['cik']
                   
                #Looking for the values and taking the average of them
                prev_value = merged_df_nan.loc[(merged_df_nan['cik'] == cik) & ((merged_df_nan['fyearq'] < fyearq) | ((merged_df_nan['fyearq'] == fyearq) & (merged_df_nan['fqtr'] < fqtr)))].sort_values(['fyearq', 'fqtr'], ascending=[False, False]).head(1)[column].item()
                next_value = merged_df_nan.loc[(merged_df_nan['cik'] == cik) & ((df['fyearq'] > fyearq) | ((merged_df_nan['fyearq'] == fyearq) & (merged_df_nan['fqtr'] > fqtr)))].sort_values(['fyearq', 'fqtr'], ascending=[True, True]).head(1)[column].item()
                average_value = np.mean([prev_value, next_value])
                merged_df_nan.loc[index, column] = average_value
                
    return merged_df_nan

优化方案（向量化操作） #

利用pandas的分组（groupby）和位移（shift）实现向量化计算，避免循环，大幅提升速度：

优化思路 #

1. 先确保数据按cik、fyearq、fqtr排序，保证周期顺序正确
2. 筛选目标列：排除指定非数值列，保留NaN数<20000的数值列
3. 按cik分组，对每个目标列计算前一期（shift(1)）和后一期（shift(-1)）的值
4. 用前/后值的均值填充NaN，同时处理边界情况（无前置/后置值时，用存在的单值填充或保留NaN）

优化代码 #

import pandas as pd
import numpy as np

def fill_nan_with_prev_next_mean(merged_df_nan):
    # 1. 按公司和周期排序，确保前后期顺序正确
    df = merged_df_nan.sort_values(by=['cik', 'fyearq', 'fqtr']).copy()
    
    # 2. 计算各列NaN数量，筛选需要填充的列
    na_counts = df.isna().sum()
    # 排除指定非数值列，保留NaN数<20000的列
    exclude_cols = ['Text', 'gvkey', 'datadate', 'prccq', 'naics', 'sic', 'processed_text', 
                   'stemmed_text', 'sentiment_score', 'fyr', 'indfmt','consol','popsrc',
                   'tic','conm','apdedateq','fdateq','rdq']
    columns_to_fill = na_counts[(na_counts < 20000) & (~na_counts.index.isin(exclude_cols))].index
    
    # 3. 按公司分组，计算前后期值的均值
    for col in columns_to_fill:
        # 分组后取前一期和后一期的值
        prev_vals = df.groupby('cik')[col].shift(1)
        next_vals = df.groupby('cik')[col].shift(-1)
        
        # 计算均值：若其中一个为NaN，则用另一个值；都为NaN则保留原NaN
        mean_vals = np.where(prev_vals.isna(), next_vals, 
                            np.where(next_vals.isna(), prev_vals, (prev_vals + next_vals)/2))
        
        # 填充原列的NaN
        df[col] = df[col].fillna(mean_vals)
    
    return df

逻辑验证方法 #

1. 抽样验证：从DataFrame中选取某几家公司的小样本，手动标记NaN位置，计算预期填充值，与代码输出对比
2. NaN数量检查：填充后统计目标列的NaN数量，确认数量减少（排除那些既无前置也无后置值的NaN）
3. 边界案例测试：测试公司第一期/最后一期的NaN，查看是否按逻辑填充（如第一期NaN只用后一期值填充）

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