主页 / coding / 问题 / 78979548
Accepted
Andi
Asked: 2024-09-13 02:45:17 +0800 CST

根据列名创建列表列

  • 772

我有一个简单的pl.DataFrame，其中有许多列仅包含布尔值。

import polars as pl

df = pl.DataFrame(
    {"s1": [True, True, False], "s2": [False, True, True], "s3": [False, False, False]}
)

shape: (3, 3)
┌───────┬───────┬───────┐
│ s1    ┆ s2    ┆ s3    │
│ ---   ┆ ---   ┆ ---   │
│ bool  ┆ bool  ┆ bool  │
╞═══════╪═══════╪═══════╡
│ true  ┆ false ┆ false │
│ true  ┆ true  ┆ false │
│ false ┆ true  ┆ false │
└───────┴───────┴───────┘

我需要添加另一个包含不同长度列表的列。任何单个行中的列表都应包含列名，其中 、 和 列的S1值s2都有s3值True。

这是我真正在寻找的：

shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    │ list         │
│ ---   ┆ ---   ┆ ---   │ ---          │
│ bool  ┆ bool  ┆ bool  │ list[str]    │
╞═══════╪═══════╪═══════╡══════════════╡
│ true  ┆ false ┆ false │ ["s1"]       │
│ true  ┆ true  ┆ false │ ["s1", "s2"] │
│ false ┆ true  ┆ false │ ["s2"]       │
└───────┴───────┴───────┴──────────────┘
python

3 个回答

  1. Best Answer

    您可以建立一个 when/then 表达式列表，然后删除空值。

    df.with_columns(
   pl.concat_list(
      pl.when(col).then(pl.lit(col)) for col in df.columns
   )
   .list.drop_nulls()
   .alias("list")
)

    shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    ┆ list         │
│ ---   ┆ ---   ┆ ---   ┆ ---          │
│ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪══════════════╡
│ true  ┆ false ┆ false ┆ ["s1"]       │
│ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴──────────────┘

    或者.list.eval()

    df.with_columns(
   pl.concat_list(pl.all())
     .list.eval(pl.lit(df.columns).gather(pl.element().arg_true()))
     .alias("list")
)

    shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    ┆ list         │
│ ---   ┆ ---   ┆ ---   ┆ ---          │
│ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪══════════════╡
│ true  ┆ false ┆ false ┆ ["s1"]       │
│ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴──────────────┘
    • 1

  2. 首先，您可以使用 创建所有列的结构struct，然后可以使用map_elements为每一行应用 lambda 函数。col如果其对应值为valTrue，lambda 函数将创建一个包含列名的列表推导。然后使用 为新列使用别名list。

    desired_output = df.with_columns(
    [
        pl.struct(df.columns)
        .map_elements(lambda x: [col for col, val in x.items() if val])
        .alias("list")
    ]
)

    或者你可以这样做

    result = (
    df.with_columns(
        [
            pl.concat_list(
                [
                    pl.when(pl.col(col)).then(pl.lit(col)).otherwise(pl.lit(None))
                    for col in df.columns
                ]
            ).alias("temp_list")
        ]
    )
    .with_columns(
        [
            pl.col("temp_list")
            .list.eval(pl.element().filter(pl.element().is_not_null()))
            .alias("list")
        ]
    )
    .drop("temp_list")
)

    打印（所需输出）

    shape: (3, 4)
┌───────┬───────┬───────┬──────────────┐
│ s1    ┆ s2    ┆ s3    ┆ list         │
│ ---   ┆ ---   ┆ ---   ┆ ---          │
│ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪══════════════╡
│ true  ┆ false ┆ false ┆ ["s1"]       │
│ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴──────────────┘
    • 0

  3. 您可以melt、filter、group_by.agg和join：

    (df.with_row_index()
   .join(df.with_row_index().melt('index').filter(pl.col('value'))
           .group_by('index').agg(pl.col('variable')),
         on='index', how='left'
        )
)

    输出：

    ┌───────┬───────┬───────┬───────┬──────────────┐
│ index ┆ s1    ┆ s2    ┆ s3    ┆ variable     │
│ ---   ┆ ---   ┆ ---   ┆ ---   ┆ ---          │
│ u32   ┆ bool  ┆ bool  ┆ bool  ┆ list[str]    │
╞═══════╪═══════╪═══════╪═══════╪══════════════╡
│ 0     ┆ true  ┆ false ┆ false ┆ ["s1"]       │
│ 1     ┆ true  ┆ true  ┆ false ┆ ["s1", "s2"] │
│ 2     ┆ false ┆ true  ┆ false ┆ ["s2"]       │
└───────┴───────┴───────┴───────┴──────────────┘
    • 0

相关问题