如何用Polars表达式计算DataFrame列的正整数序列加权和？

要实现按1、2、3…连续正整数作为权重的列加权和，直接用Polars的向量化表达式就能高效完成，无需循环，代码如下：

import polars as pl

df = pl.DataFrame({"a": [2, 4, 2, 1, 2, 1, 3, 6, 7, 5]})

# 计算加权和并提取结果
weighted_sum = df.select(
    pl.sum(pl.col("a") * pl.int_range(start=1, end=pl.count() + 1)).alias("weighted_sum")
).item()

print(weighted_sum)  # 输出：218

核心逻辑说明 #

• 生成权重序列：pl.int_range(start=1, end=pl.count() + 1)会生成从1开始、长度与DataFrame行数完全匹配的连续整数序列。pl.count()动态获取当前DataFrame的总行数，确保权重序列长度随数据行数自动调整。
• 向量化加权计算：将目标列a与权重序列逐元素相乘，Polars会以向量化方式完成此运算，性能远高于Python循环。
• 求和得到结果：用pl.sum()对所有加权后的元素求和，最后通过.item()提取单个数值结果。

这种方法完全基于Polars原生表达式，支持Lazy API（只需将df.select改为df.lazy().select().collect()），在处理百万级以上行数的大型DataFrame时依然能保持高效运算。

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