如何在Python Polars中拆分列表列并按照比例分割整数列？

Question

我有一个数据集，里面包含了物品类型和每种物品的数量。不过，有些行的物品类型这一列里不是单一类型，而是包含了多个类型的列表。我想把这些类型的列表拆分成单独的行，并且按比例分配数量。

在使用pandas时，我通常会先拆分（explode），然后根据索引进行分组，类似于这里描述的过程（那页上的示例数据和我正在处理的数据类似）。

我可以在polars数据框中创建自己的索引，或者在内存能容纳的情况下，转换到pandas再转换回来，但有没有更好的方法在Polars中做到这一点呢？

编辑：

评论中似乎有些混淆，为了更清楚，我的数据格式是这样的：

┌───────────────────┬─────┐
│ cat               ┆ qty │
│ ---               ┆ --- │
│ list[str]         ┆ i64 │
╞═══════════════════╪═════╡
│ ["green", "blue"] ┆ 23  │
│ ["green"]         ┆ 23  │
│ ["red"]           ┆ 4   │
│ ["blue"]          ┆ 5   │
│ ["red", "blue"]   ┆ 15  │
└───────────────────┴─────┘

我想要的是：

┌───────┬──────┐
│ cat   ┆ qty  │
│ ---   ┆ ---  │
│ str   ┆ f64  │
╞═══════╪══════╡
│ green ┆ 11.5 │
│ blue  ┆ 11.5 │
│ green ┆ 23.0 │
│ red   ┆ 4.0  │
│ blue  ┆ 5.0  │
│ red   ┆ 7.5  │
│ blue  ┆ 7.5  │
└───────┴──────┘

我可以通过以下方式获得：

#convert to pandas
df_out = df.to_pandas()
#exploding cats
df_out = df_out.explode('cat')
#correct for multiplied quantity column
df_out['qty'] /= df_out['qty'].groupby(level=0).transform('count')
#back to polars
df_out = pl.DataFrame(df_out)
#print
print(df_out)

也可以用类似于这个的脚本来实现（感谢@Henricks找到那篇文章！）。

但我想知道有没有更好的方法？理想情况下，希望能在懒加载模式下完成，而不需要创建一个新的列作为伪索引？不过我还没有想到什么有效的办法。谢谢！

Answer 1

和pandas的方法类似，你可以把qty除以对应的cat列表中的元素数量。

(
    df
    .with_columns(pl.col("qty") / pl.col("cat").list.len())
    .explode("cat")
)

shape: (7, 2)
┌───────┬──────┐
│ cat   ┆ qty  │
│ ---   ┆ ---  │
│ str   ┆ f64  │
╞═══════╪══════╡
│ green ┆ 11.5 │
│ blue  ┆ 11.5 │
│ green ┆ 23.0 │
│ red   ┆ 4.0  │
│ blue  ┆ 5.0  │
│ red   ┆ 7.5  │
│ blue  ┆ 7.5  │
└───────┴──────┘

注意。 和pandas不同的是，我们在“展开”列表列之前就进行这个操作（这样我们仍然可以获取列表的长度）。