英文:
Weighted sum of a column in Polars dataframe
问题
我有一个 Polars 数据帧，我想计算特定列的加权和，权重仅为正整数序列，例如1、2、3、...。
例如，假设我有以下数据帧。
import polars as pl
df = pl.DataFrame({"a": [2, 4, 2, 1, 2, 1, 3, 6, 7, 5]})
我想要的结果是
218 (= 2*1 + 4*2 + 2*3 + 1*4 + ... + 7*9 + 5*10)
如何通过仅使用通用 Polars 表达式来实现这一目标？（我之所以要仅使用 Polars 表达式来解决问题，是出于性能考虑）
注意：这个示例只是一个简单的示例，其中只有10个数字，但一般情况下，数据帧的高度可以是任何正整数。
感谢您的帮助。
英文:
I have a Polars dataframe and I want to calculate a weighted sum of a particular column and the weights is just the positive integer sequence, e.g., 1, 2, 3, ...
For example, assume I have the following dataframe.
import polars as pl
df = pl.DataFrame({"a": [2, 4, 2, 1, 2, 1, 3, 6, 7, 5]})
The result I want is
218 (= 2*1 + 4*2 + 2*3 + 1*4 + ... + 7*9 + 5*10)
How can I achieve this by using only general polars expressions? (The reason I want to use just polars expressions to solve the problem is for speed considerations)
Note: The example is just a simple example where there are just 10 numbers there, but in general, the dataframe height can be any positive number.
Thanks for your help..
答案1
得分: 1
这种加权总和可以使用点积（.dot() 方法）计算。要生成从1到n的范围（权重），可以使用 pl.arange(1, n+1)。
如果只需要计算加权总和的结果：
df.select(
pl.col("a").dot(pl.arange(1, pl.count()+1))
) #.item() - to get value (218)
保持数据框：
df.with_columns(
pl.col("a").dot(pl.arange(1, pl.count()+1)).alias("weighted_sum")
)
┌─────┬──────────────┐
│ a ┆ weighted_sum │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪══════════════╡
│ 2 ┆ 218 │
│ 4 ┆ 218 │
│ ... ┆ ... │
│ 3 ┆ 218 │
│ 5 ┆ 218 │
└─────┴──────────────┘
在 groupby 上下文中：
df.groupby("some_cat_col", maintain_order=True).agg([
pl.col("a").dot(pl.arange(1, pl.count()+1))
])
英文:
Such weighted sum can be calculated using dot product (.dot() method). To generate range (weights) from 1 to n, you can use pl.arange(1, n+1).
If you just need to calculate result of weighted sum:
df.select(
pl.col("a").dot(pl.arange(1, pl.count()+1))
) #.item() - to get value (218)
Keep dataframe
df.with_columns(
pl.col("a").dot(pl.arange(1, pl.count()+1)).alias("weighted_sum")
)
┌─────┬──────────────┐
│ a ┆ weighted_sum │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪══════════════╡
│ 2 ┆ 218 │
│ 4 ┆ 218 │
│ ... ┆ ... │
│ 3 ┆ 218 │
│ 5 ┆ 218 │
└─────┴──────────────┘
In groupby context
df.groupby("some_cat_col", maintain_order=True).agg([
pl.col("a").dot(pl.arange(1, pl.count()+1))
])
答案2
得分: 0
你可以使用以下代码计算带有索引+1的系列a的点积：
import polars as pl
df = pl.DataFrame({"a": [2, 4, 2, 1, 2, 1, 3, 6, 7, 5]})
print(df["a"].dot(df.index + 1))
或者，你也可以使用__matmul__运算符 @，如下：
print(df["a"] @ (df.index + 1))
英文:
You should be able to dot the series a with the index + 1
import polars as pl
df = pl.DataFrame({"a": [2, 4, 2, 1, 2, 1, 3, 6, 7, 5]})
print(df["a"].dot(df.index+1))
Alternatively, you can use the __matmul__ operator @
print(df["a"] @ (df.index + 1))
通过集体智慧和协作来改善编程学习和解决问题的方式。致力于成为全球开发者共同参与的知识库，让每个人都能够通过互相帮助和分享经验来进步。
评论