Python 基于另一个数据帧的一个数据帧的列求和

发布于07月05日

我有两个数据帧，看起来像这样:

df1 = pl.DataFrame(
    {
        "Name": ["A", "B", "C", "D"],
        "Year": [2001, 2003, 2003, 2004]
    }
)

df2 = pl.DataFrame(
    {
        "Name": ["A", "B", "C", "D"],
        "2001": [111, 112, 113, 114],
        "2002": [221, 222, 223, 224],
        "2003": [331, 332, 333, 334],
        "2004": [441, 442, 443, 444]
    }
)

我想对第二个df(Df2)的每一年列求和，只考虑df1中相应年份为同年或更晚的名字.所需输出:

┌──────┬──────┐
│ Year ┆  Sum │
╞══════╪══════╡
│ 2001 ┆  111 │  
│ 2002 ┆  221 │
│ 2003 ┆  996 │  (= 331 + 332 + 333)
│ 2004 ┆ 1770 │  (= 441 + 442 + 443 + 444)
└──────┴──────┘

我刚接触北极星(来自Pandas )，我不确定该怎么做.任何帮助都将不胜感激.

推荐答案

你可以用一种与df1更"兼容"的形式来表达它.它将指定的列标题"取消透视"为行.在这一点上，Year应该从str变为int列:

df2.melt(id_vars='Name', variable_name='Year').with_columns(
    pl.col('Year').cast(pl.Int64)
)

shape: (16, 3)
┌──────┬──────┬───────┐
│ Name ┆ Year ┆ value │
│ ---  ┆ ---  ┆ ---   │
│ str  ┆ i64  ┆ i64   │
╞══════╪══════╪═══════╡
│ A    ┆ 2001 ┆ 111   │
│ B    ┆ 2001 ┆ 112   │
│ C    ┆ 2001 ┆ 113   │
│ D    ┆ 2001 ┆ 114   │
│ A    ┆ 2002 ┆ 221   │
│ B    ┆ 2002 ┆ 222   │
│ C    ┆ 2002 ┆ 223   │
│ D    ┆ 2002 ┆ 224   │
│ A    ┆ 2003 ┆ 331   │
│ B    ┆ 2003 ┆ 332   │
│ C    ┆ 2003 ┆ 333   │
│ D    ┆ 2003 ┆ 334   │
│ A    ┆ 2004 ┆ 441   │
│ B    ┆ 2004 ┆ 442   │
│ C    ┆ 2004 ┆ 443   │
│ D    ┆ 2004 ┆ 444   │
└──────┴──────┴───────┘

然后，您可以使用Name上的df1对其进行join，并根据您指定的年份条件筛选结果行:

(...).join(df1, on='Name').filter(pl.col('Year') >= pl.col('Year_right'))

从那里开始，Year号公路上的groupby和agg就是一个简单的累加.总而言之:

df2.melt(id_vars='Name', variable_name='Year').with_columns(
    pl.col('Year').cast(pl.Int64)
).join(df1, on='Name').filter(pl.col('Year') >= pl.col('Year_right')).groupby(
    'Year', maintain_order=True
).agg(
    Sum=pl.col('value').sum()
)

shape: (4, 2)
┌──────┬──────┐
│ Year ┆ Sum  │
│ ---  ┆ ---  │
│ i64  ┆ i64  │
╞══════╪══════╡
│ 2001 ┆ 111  │
│ 2002 ┆ 221  │
│ 2003 ┆ 996  │
│ 2004 ┆ 1770 │
└──────┴──────┘

(maintain_order是可选的，但速度较慢，如果确实需要，可以达到与agg之后的.sort('Year')相同的效果.)