分组计数值作为列

Question

我想按一些列分组，但其中一列是'Fire'或'Water'。我需要在单独的列中计算'Fire'和'Water'的出现次数，并且还需要一个'total'列来计算'Fire'和'Water'的总和

的例子:

df = pl.DataFrame({'ID': [0,0,1, 1, 1,0], 'Type': ['Fire', 'Fire', 'Fire', 'Water', 'Water', 'Water'], })
>>>df 
shape: (6, 2)
┌─────┬───────┐
│ ID  ┆ Tye
│ i64 ┆ str   │
│ 0   ┆ Fire  │
│ 0   ┆ Fire  │
│ 1   ┆ Fire  │
│ 1   ┆ Water │
│ 1   ┆ Water │
│ 0   ┆ Water │
df = df.groupby(['ID', 'Type']).agg(pl.count())
shape: (4, 3)
┌─────┬───────┬───────┐
│ ID  ┆ Type  ┆ count │
│ --- ┆ ---   ┆ ---   │
│ i64 ┆ str   ┆ u32   │
╞═════╪═══════╪═══════╡
│ 1   ┆ Fire  ┆ 1     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 1   ┆ Water ┆ 2     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 0   ┆ Fire  ┆ 2     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 0   ┆ Water ┆ 1     │
└─────┴───────┴───────┘

我最想拥有的是:

┌─────┬───────┬──────┬───────┐
│ ID  ┆ water ┆ fire ┆ total │
│ --- ┆ ---   ┆ ---  ┆ ---   │
│ i64 ┆ u32   ┆ u32  ┆ u32   │
╞═════╪═══════╪══════╪═══════╡
│ 1   ┆ 2     ┆ 1    ┆ 3     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 0   ┆ 1     ┆ 2    ┆ 3     │
└─────┴───────┴──────┴───────┘

我该怎么做呢?

Answer 1

我们可以在DataFrame上使用pivot函数将DataFrame从"长格式"转换为"长格式"。到"宽格式">

df = (
df
.pivot(
values=['Type'],
index=['ID'],
columns=['Type'],
aggregate_fn='count'
)
)
df

shape: (2, 3)
┌─────┬──────┬───────┐
│ ID  ┆ Fire ┆ Water │
│ --- ┆ ---  ┆ ---   │
│ i64 ┆ u32  ┆ u32   │
╞═════╪══════╪═══════╡
│ 0   ┆ 2    ┆ 1     │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 1   ┆ 1    ┆ 2     │
└─────┴──────┴───────┘

从这里开始，我们可以使用polars.sum表达式创建总行。当与表达式列表一起使用时，polars.sum表达式将横向地跨行求和。

df.with_column(
pl.sum([col_nm for col_nm
in df.columns
if col_nm not in ["ID"]]).alias('total')
)

shape: (2, 4)
┌─────┬──────┬───────┬───────┐
│ ID  ┆ Fire ┆ Water ┆ total │
│ --- ┆ ---  ┆ ---   ┆ ---   │
│ i64 ┆ u32  ┆ u32   ┆ u32   │
╞═════╪══════╪═══════╪═══════╡
│ 0   ┆ 2    ┆ 1     ┆ 3     │
├╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 1   ┆ 1    ┆ 2     ┆ 3     │
└─────┴──────┴───────┴───────┘

这个堆栈溢出问题解释了更多关于melt和pivot如何一起工作以在"长格式"之间转换的信息。和"宽格式";DataFrames .

注意:pivot只适用于dataframe，不适用LazyFrames