Pandas按行计算符合条件的首个值及对应权重结果的高效实现

解法思路 #

直接用numpy向量化操作实现，完全避免行遍历，效率远高于apply逐行计算，数据量越大优势越明显，同时也支持后续扩展f列数量的场景。

实现代码 #

import pandas as pd
import numpy as np

# 原始数据构造
df = pd.DataFrame(
    [
        [10,20,40],
        [2,1,26],
        [1, 2, 60],
    ], columns = ['f1', 'f2', 'f3']
)
df['cumsum'] = df.sum(axis=1)
df['cumsum_perc'] = (df['cumsum'] * 0.1).astype(int)

# 固定配置：f列顺序、对应权重
f_cols = ['f1', 'f2', 'f3']
pon = np.array([1, 2, 3])

# 提取数组做向量化计算
f_vals = df[f_cols].values
T = df['cumsum_perc'].values.reshape(-1, 1)

# 计算每行f_le的索引：没有符合条件的列时默认取第一个f列
mask = f_vals <= T
temp = np.where(mask, f_vals, -np.inf)
le_idx = temp.argmax(axis=1)

# 取对应值计算余数
w = pon[le_idx]
v_le = f_vals[np.arange(len(df)), le_idx]
rem = np.maximum(T.flatten() - v_le, 0)

# 计算余数占比：最右列无后续列时占比直接取0
next_vals = np.pad(f_vals, ((0,0), (0,1)), constant_values=1)[:, 1:]
rem_perc = np.where(le_idx == len(f_cols)-1, 0, rem / next_vals[np.arange(len(df)), le_idx])

# 赋值最终结果
df['result'] = w + rem_perc

结果验证 #

运行后得到的result列和示例计算结果完全匹配：

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