You need to enable JavaScript to run this app.
优惠活动
大模型
产品
解决方案
定价
更多

简化加权移动平均计算:移除.v.to_numpy()直接获取NumPy数组的可行性

简化加权移动平均计算:移除.v.to_numpy()直接获取NumPy数组的可行性

当然可以简化这段代码!你遇到的问题是因为对多列DataFrame调用rolling.apply时,pandas会把函数返回的标量广播到所有列,导致结果是一个所有列内容重复的DataFrame。这里有几个实用的方法可以直接拿到NumPy数组,省去.v.to_numpy()的步骤:

方法1:直接提取结果数组的第一列(不依赖列名)

这种方法最简单,不需要修改原有逻辑,只是把依赖列名的.v.to_numpy()换成更通用的数组索引方式,代码更健壮(就算以后列名变更也不用调整):

accs = (pd.DataFrame({'v': values[start_index:], 'w': weights[start_index:]})
        .rolling(window=window_len, min_periods=1, method='table')
        .apply(weighted_mean, engine='numba', raw=True)
        .to_numpy()[:, 0][a:b])

.to_numpy()把DataFrame转成二维NumPy数组,[:, 0]取第一列的结果(因为所有列都相同,取任意一列都可以),最后再执行切片[a:b]。

方法2:提前打包数据,让rolling直接返回Series

如果想从根源避免生成多列DataFrame,可以把v和w打包成单个Series的元素,这样rolling.apply的结果会是一个Series而非DataFrame,直接转NumPy数组即可:

# 把v和w组合成二维数组,再转成Series(每个元素是[v, w]的数组)
combined_data = pd.Series(np.column_stack([values[start_index:], weights[start_index:]]).tolist())

accs = (combined_data.rolling(window=window_len, min_periods=1, method='table')
        .apply(lambda window: weighted_mean(np.vstack(window)), engine='numba', raw=True)
        .to_numpy()[a:b])

这里np.vstack(window)会把窗口内的[v,w]元素重新拼成(window_len, 2)的二维数组,和原来weighted_mean接收的输入格式完全一致,不需要修改你的加权平均函数。

额外提示:关于numba引擎的小细节

因为你用了engine='numba'和raw=True,这已经是效率很高的写法——numba会编译你的weighted_mean函数,避免pandas对象的额外开销。上面两种方法都能保留这个效率优势,不会影响计算速度。

备注:内容来源于stack exchange,提问作者clime

相关产品推荐
方舟 Agent Plan

超全模态模型 × Harness 升级,最新支持 Deepseek-V4.1-Flash、GLM-5.3 系列、Doubao-Seedream-5.0-pro、Kimi-K3 (部分), 限时 9.9 元起

最近更新时间:2026.04.22 15:49:55