Python 极点替换值大于组内另一个极点数据帧的最大值

发布于03月16日

我有两个DataFrame:

import polars as pl

df1 = pl.DataFrame(
    {
        "group": ["A", "A", "A", "B", "B", "B"],
        "index": [1, 3, 5, 1, 3, 8],
    }
)

df2 = pl.DataFrame(
    {
        "group": ["A", "A", "A", "B", "B", "B"],
        "index": [3, 4, 7, 2, 7, 10],
    }
)

我想用df1中每组的largest index来给df2中的index加上限.两个DataFrame中的组是相同的.

预期输出df2:

shape: (6, 2)
┌───────┬───────┐
│ group ┆ index │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ A     ┆ 3     │
│ A     ┆ 4     │
│ A     ┆ 5     │
│ B     ┆ 2     │
│ B     ┆ 7     │
│ B     ┆ 8     │
└───────┴───────┘

推荐答案

您可以通过df1计算每个组的最大值，然后计算clip df2:

out = df2.with_columns(
    pl.col('index').clip(
        upper_bound=df1.select(pl.col('index').max().over('group'))['index']
    )
)

输出:

shape: (6, 2)
┌───────┬───────┐
│ group ┆ index │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ A     ┆ 3     │
│ A     ┆ 4     │
│ A     ┆ 5     │
│ B     ┆ 2     │
│ B     ┆ 7     │
│ B     ┆ 8     │
└───────┴───────┘

或者，如果两个组在两个框架中不一定相同，则可以将group_by.max与join对齐:

df1 = pl.DataFrame(
    {
        "group": ["A", "A", "A", "B", "B", "B"],
        "index": [1, 3, 5, 1, 3, 7],
    }
)

df2 = pl.DataFrame(
    {
        "group": ["A", "A", "A", "B", "B", "B", "B"],
        "index": [3, 4, 7, 2, 7, 8, 9],
    }
)

out = df2.with_columns(
    pl.col('index').clip(
        upper_bound=df2.join(df1.group_by('group').max(), on='group')['index_right']
    )
)

输出:

shape: (7, 2)
┌───────┬───────┐
│ group ┆ index │
│ ---   ┆ ---   │
│ str   ┆ i64   │
╞═══════╪═══════╡
│ A     ┆ 3     │
│ A     ┆ 4     │
│ A     ┆ 5     │
│ B     ┆ 2     │
│ B     ┆ 7     │
│ B     ┆ 7     │
│ B     ┆ 7     │
└───────┴───────┘