主页 / coding / 问题 / 78999687
Accepted
rindis
Asked: 2024-09-19 02:22:22 +0800 CST

极坐标使所有组的大小相同

  • 772

问题

我正在尝试使给定数据框的所有组都具有相同的大小。在下面的起点中，我展示了一个我想要转换的数据框示例。在目标中，我尝试演示我想要实现的目标。我想按列分组group，使所有组的大小为4，并用填充“缺失”值null- 我希望这很清楚。

我尝试过几种方法，但还是没能解决这个问题。

起点

dfa = pl.DataFrame(data={'group': ['a', 'a', 'a', 'b', 'b', 'c'],
                         'value': ['a1', 'a2', 'a3', 'b1', 'b2', 'c1']})
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ c     ┆ c1    │
└───────┴───────┘

目标

>>> make_groups_uniform(dfa, group_by='group', group_size=4)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ a     ┆ null  │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ b     ┆ null  │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘

软件包版本

polars: 1.1.0

解决方案

我根据下面@jqurious 的回答得出了这个结论。

>>> import polars as pl

>>> dfa = pl.DataFrame(data={'group': ['a', 'a', 'a', 'b', 'b', 'c'],
...                          'value': ['a1', 'a2', 'a3', 'b1', 'b2', 'c1']})

┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ c     ┆ c1    │
└───────┴───────┘

>>> (dfa
...  .with_columns(group_size=pl.col('group')
...                             .count()
...                             .over('group')
...                             .max()
...                             .explode())
...  .group_by('group', maintain_order=True)
...  .agg(pl.all()
...         .append(pl.repeat(None, pl.col('group_size') - pl.len())))
...  .select(pl.exclude('group_size'))
...  .explode(pl.exclude('group')))

┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘
python

4 个回答

  1. Best Answer

    您可以使用它pl.repeat()来生成空值和.append()它们。

    group_size = 4

(df.group_by("group", maintain_order=True)
   .agg(pl.all().append(pl.repeat(None, group_size - pl.len().cast(int))))
   .explode(pl.exclude("group"))
) 

    shape: (12, 2)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ a     ┆ null  │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ b     ┆ null  │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘

    ☣ 警告 ☣

    pl.len()目前是一个无符号整数（将来可能会改变）

    我们转换为整数以避免潜在的溢出问题。

    df.group_by("group").agg(int = 2 - pl.len().cast(int), uint32 = 2 - pl.len())

    shape: (3, 3)
┌───────┬─────┬────────────┐
│ group ┆ int ┆ uint32     │
│ ---   ┆ --- ┆ ---        │
│ str   ┆ i64 ┆ u32        │
╞═══════╪═════╪════════════╡
│ a     ┆ -1  ┆ 4294967295 │ # <- pl.repeat(None, 4294967295) would be bad.
│ b     ┆ 0   ┆ 0          │
│ c     ┆ 1   ┆ 1          │
└───────┴─────┴────────────┘
    1. https://github.com/pola-rs/polars/issues/17722
    • 5

  2. 您可以使用 找到最大的组.group_by().len().max()，然后 groupby 并将组中的每个值转换为结构体，使用'max_width'策略解除结构体的嵌套，然后融化宽数据。之后，您可以填补最大组和所需组大小之间的差异。

    def make_groups_uniform(df: pl.DataFrame, group_by: str, value_col: str, group_size: int):
    gb = df.group_by(group_by)
    largest_group_size = gb.len().max()['len'].item()
    out = gb.all().with_columns(
        pl.col(value_col)
          .list
          .to_struct(n_field_strategy='max_width')
    ).unnest(
         value_col
    ).melt(group_by, value_name=value_col)[[group_by, value_col]]
    n = group_size - largest_group_size

    if n > 0:
        keys = df[group_by].unique().to_list()
        pad = pl.DataFrame({group_by: keys * n, value_col: ['null']*n*len(keys)})
        out = pl.concat([out, pad])
    return out.sort(group_by)

    运行代码：

    make_groups_uniform(dfa, 'group', 'value', 4)
# returns:
shape: (12, 2)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ a     ┆ null  │
│ b     ┆ b1    │
│ …     ┆ …     │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘
    • 1

  3. 你可以通过groupby按组列对 df 进行分组来解决这个问题。然后，NULL如果组 <= 4 ，则用值填充每个组，然后我们在处理后连接所有组pl.concat.

    这里有一个代码片段可以指导您。

    ef make_groups_uniform(df, group_by, group_size):
    def padding_each_group(group: pl.DataFrame) -> pl.DataFrame:
        missing_data_count = group_size - group.shape[0]
        if missing_data_count > 0:
            # missing values
            missing_df = pl.DataFrame(
                {
                    group_by: [group[group_by][0]]
                    * missing_data_count,  # Fill with the group value
                    "value": [None] * missing_data_count,  # Fill with None
                }
            )
            # Concatenate the original group with the missing DataFrame
            group = pl.concat([group, missing_df])
        return group

    # Collect each group as separate DataFrames
    groups = []
    for _, df_group in df.group_by(group_by, maintain_order=True):
        padded_group = padding_each_group(df_group)
        groups.append(padded_group)

    # concat all groups
    output = pl.concat(groups)

    return output

    打印（输出）

    shape: (12, 2)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ a     ┆ null  │
│ b     ┆ b1    │
│ …     ┆ …     │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘```
    • 1

  4. 这种方法的优点是我们实际上不需要转换原始 DataFrame（除了如果您想重新排列组，可能需要排序），我们只需要创建额外的行并将它们附加到原始 DataFrame 中。

    我对答案进行了一些调整，基于您希望组的大小为所有组大小的最大值的假设，但它对于固定的大小也有效group_size。

    # you can use fixed group_size instead of pl.col.len.max() as well

pl.concat([
    dfa,
    (
        dfa.group_by("group").len()
        .select(
            pl.col.group.repeat_by(pl.col.len.max() - pl.col.len).explode(),
            value = pl.lit(None)
        )
    )
]).sort("group"))

    shape: (10, 2)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ a     ┆ null  │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘

    如果您需要固定的组大小，那么repeat()性能可能会更高，但想法是一样的 - 只生成额外的行并将它们附加回原始 DataFrame。

    group_size = 3

# you can make it dynamic as well though
# group_size = dfa.group_by("group").len().max()["len"].item()

pl.concat([
    dfa,
    (
        dfa.group_by("group")
        .agg(value = pl.repeat(None, group_size - pl.len().cast(int)))
        .explode("value")
    )
]).sort("group")

    shape: (10, 2)
┌───────┬───────┐
│ group ┆ value │
│ ---   ┆ ---   │
│ str   ┆ str   │
╞═══════╪═══════╡
│ a     ┆ a1    │
│ a     ┆ a2    │
│ a     ┆ a3    │
│ a     ┆ null  │
│ b     ┆ b1    │
│ b     ┆ b2    │
│ b     ┆ null  │
│ c     ┆ c1    │
│ c     ┆ null  │
│ c     ┆ null  │
└───────┴───────┘
    • 0

相关问题