python-polars - 部分基于其他列名称创建新列

Question

总的来说，我对 Polars 和 Python 都很陌生。我有一个有点不寻常的问题，需要一些帮助。我有一个包含 50 多个 0/1 列的数据框。我需要创建一个新列，其中包含每列的逗号分隔列表，其中包含 1 但使用部分列名。如果 hccx = 1，则将 x 附加到字符串列。一个简化的例子:

df=pl.DataFrame(
    {'id':[1,2,3], 'hcc1':[0,1,1],'hcc2':[0,0,1],'hcc5':[0,1,1],'hcc8':[1,0,0]}
)

shape: (3, 5)
┌─────┬──────┬──────┬──────┬──────┐
│ id  ┆ hcc1 ┆ hcc2 ┆ hcc5 ┆ hcc8 │
│ --- ┆ ---  ┆ ---  ┆ ---  ┆ ---  │
│ i64 ┆ i64  ┆ i64  ┆ i64  ┆ i64  │
╞═════╪══════╪══════╪══════╪══════╡
│ 1   ┆ 0    ┆ 0    ┆ 0    ┆ 1    │
│ 2   ┆ 1    ┆ 0    ┆ 1    ┆ 0    │
│ 3   ┆ 1    ┆ 1    ┆ 1    ┆ 0    │
└─────┴──────┴──────┴──────┴──────┘

我想创建一个新列(字符串类型)hccall，如下所示:

<表类=“s-表”><标题>idhccall <正文>1821,531,2,5

我想象某种类型的列表理解在以“hcc”开头的列上循环会起作用，但我有点卡住了。我可以创建一个循环，但不确定如何从循环内附加到列。有什么巧妙的想法吗？

Answer 1

我认为最简单的选择是首先融化你的数据帧，这样你的数据帧中每列每个 id 一行。然后您可以过滤等于 1 的行。然后您可以将原始行的最后一个字母聚合到一个列表中。使用list.join，您可以将列表组合成一个以逗号分隔的字符串

(
    df.melt(id_vars="id")
    .filter(pl.col("value") == 1)
    .groupby("id")
    .agg(pl.col("variable").str.slice(3))
    .with_columns(pl.col("variable").list.join(","))
)

shape: (3, 2)
┌─────┬──────────┐
│ id  ┆ variable │
│ --- ┆ ---      │
│ i64 ┆ str      │
╞═════╪══════════╡
│ 2   ┆ 1,5      │
│ 3   ┆ 1,2,5    │
│ 1   ┆ 8        │
└─────┴──────────┘