如何优化百万级DataFrame迭代的k_nn_averaging算法？

你的代码性能瓶颈核心在于逐行循环处理百万级DataFrame——pandas的iloc/loc在循环中反复调用会带来极大的索引开销，加上每次循环里的列表操作、随机函数调用，直接导致整体效率极低。以下是几个可落地的优化方案：

1. 完全向量化操作，彻底移除循环 #

核心思路是用numpy批量生成所有需要的随机邻居索引和因子，通过数组运算完成全量计算，彻底避免逐行处理：

import numpy as np
import pandas as pd

def optimized_k_nn_averaging(df: pd.DataFrame, k: int = 15, use_abs_value: bool = False) -> pd.DataFrame:
    mod_cols = helper.modifiable_columns
    # 提取可修改列转为numpy数组，减少pandas索引开销
    df_data = df[mod_cols].astype(float).values
    n_rows = df_data.shape[0]
    
    # 批量计算每行的邻居范围边界
    left_bounds = np.maximum(0, np.arange(n_rows) - k)
    right_bounds = np.minimum(n_rows, np.arange(n_rows) + k + 1)  # range左闭右开，故+1
    
    # 批量生成随机邻居索引，避免选中自身
    neighbor_indices = np.random.randint(left_bounds, right_bounds, size=n_rows)
    # 修正选中自身的情况：直接切换到相邻索引，无需重试
    self_mask = neighbor_indices == np.arange(n_rows)
    neighbor_indices[self_mask] = np.where(
        left_bounds[self_mask] < np.arange(n_rows)[self_mask],
        np.arange(n_rows)[self_mask] - 1,
        np.arange(n_rows)[self_mask] + 1
    )
    
    # 批量生成随机因子，转成列向量支持广播运算
    factors = np.random.uniform(0, 1, size=n_rows)[:, np.newaxis]
    
    # 批量计算差值并更新
    diff = df_data - df_data[neighbor_indices]
    if use_abs_value:
        diff = np.abs(diff)
    updated_data = df_data + factors * diff
    
    # 生成结果DataFrame
    df_averaged = df.copy()
    df_averaged[mod_cols] = updated_data
    return df_averaged

2. 关键优化点说明 #

• 移除Python循环：numpy底层是C实现的批量数组操作，比Python逐行循环效率高100~1000倍，百万行数据能从“小时级”压缩到“秒级”。
• 批量生成随机值：一次性生成所有邻居索引和因子，避免在循环中反复调用choice/uniform带来的函数调用开销。
• 用numpy数组替代Series：直接操作numpy数组跳过pandas的索引层，大幅减少数据访问耗时。
• 简化类型转换：仅对可修改列做一次类型转换，避免重复操作浪费资源。

3. 额外小优化 #

• 如果精度允许，可将astype(float)改为astype(np.float32)，减少内存占用的同时提升运算速度。
• 若DataFrame内存紧张，可考虑分块处理（比如每10万行处理一次），避免内存溢出。

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