Pandas按每行Top x指定数量取最高周列值计算均值的实现方法

直接采用逐行自定义计算的方式即可，逻辑清晰易维护，普通数据集场景下效率完全够用：

import pandas as pd
import numpy as np

# 构造示例DataFrame
df = pd.DataFrame({'week 1' : [1.0, 5.0, 7.0, 6.0, np.nan],
               'week 2' : [3.0, np.nan, 9.0, 8.0, np.nan],
               'week 3' : [1.0, 2.0, 2.0, 1.0, 6.0],
               'week 4' : [np.nan, 4.0, 2.0, 7.0, 6.0],
               'Top x' : [3.0, 2.0, 4.0, 3.0, 3.0]})

def calc_row_mean(row):
    # 提取当前行所有week列的非空值
    week_valid = row.filter(like='week').dropna()
    x = int(row['Top x'])
    # 实际计算的数量取x和有效数量的最小值，兼容有效数不足x的情况
    calc_num = min(x, len(week_valid))
    if calc_num == 0:
        return np.nan
    # 取最大的calc_num个值求均值
    return week_valid.nlargest(calc_num).mean()

df['Mean'] = df.apply(calc_row_mean, axis=1)

执行后得到的Mean列和预期结果完全一致。如果是处理百万级以上的超大数据集，可以再优化为向量化操作进一步提升运行速度。

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