Значения подсчета группировки как новые столбцы ⇐ Python

[Anonymous]

Я хочу сгруппировать несколько столбцов, но один из столбцов — «Огонь» или «Вода».
Мне нужно подсчитать количество случаев «Огонь» и «Вода» в отдельных столбцах, а также иметь столбец «Итог», в котором подсчитывается сумма «Огонь» и «Вода».
Пример:

Код: Выделить всё

import polars as pl

df = pl.DataFrame({
'ID': [0, 0,1, 1, 1, 0],
'Type': ['Fire', 'Fire', 'Fire', 'Water', 'Water', 'Water'],
})

Код: Выделить всё

shape: (6, 2)
┌─────┬───────┐
│ ID  ┆ Type  │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 0   ┆ Fire  │
│ 0   ┆ Fire  │
│ 1   ┆ Fire  │
│ 1   ┆ Water │
│ 1   ┆ Water │
│ 0   ┆ Water │
└─────┴───────┘

Мой текущий подход:

Код: Выделить всё

df.group_by('ID', 'Type').agg(count = pl.col("Type").count())

Код: Выделить всё

shape: (4, 3)
┌─────┬───────┬───────┐
│ ID  ┆ Type  ┆ count │
│ --- ┆ ---   ┆ ---   │
│ i64 ┆ str   ┆ u32   │
╞═════╪═══════╪═══════╡
│ 1   ┆ Water ┆ 2     │
│ 1   ┆ Fire  ┆ 1     │
│ 0   ┆ Water ┆ 1     │
│ 0   ┆ Fire  ┆ 2     │
└─────┴───────┴───────┘

Мне бы хотелось:

Код: Выделить всё

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

Как бы я это сделал?

Подробнее здесь: https://stackoverflow.com/questions/729 ... ew-columns