主页 / user-4436517

rindis's questions

Martin Hope
rindis
Asked: 2024-09-19 02:22:22 +0800 CST
  • 8

问题

我正在尝试使给定数据框的所有组都具有相同的大小。在下面的起点中，我展示了一个我想要转换的数据框示例。在目标中，我尝试演示我想要实现的目标。我想按列分组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