如何在 Polars 中按百分比差异对项目进行分组？

Question

我想对值进行分组，以便每个组项目之间的差异保持在一定的百分比内。例如。每当一个项目超过第一个组元素 5% 时，它就会进入新组。作为回报，我需要第一个组值。 5% 阈值的示例，其中给出“a”，必须计算“group”和“groupFirst”：

import polars as pl

df = pl.DataFrame({'a': [100, 103, 105, 106, 105, 104, 103, 106, 100, 102],
    'group': [0, 0, 1, 1, 1, 1, 1, 1, 2, 2], 
    'groupFirst': [100, 100, 105, 105, 105, 105, 105, 105, 100, 100]})

print(df)

shape: (10, 3)
┌─────┬───────┬────────────┐
│ a   ┆ group ┆ groupFirst │
│ --- ┆ ---   ┆ ---        │
│ i64 ┆ i64   ┆ i64        │
╞═════╪═══════╪════════════╡
│ 100 ┆ 0     ┆ 100        │
│ 103 ┆ 0     ┆ 100        │
│ 105 ┆ 1     ┆ 105        │
│ 106 ┆ 1     ┆ 105        │
│ 105 ┆ 1     ┆ 105        │
│ 104 ┆ 1     ┆ 105        │
│ 103 ┆ 1     ┆ 105        │
│ 106 ┆ 1     ┆ 105        │
│ 100 ┆ 2     ┆ 100        │
│ 102 ┆ 2     ┆ 100        │
└─────┴───────┴────────────┘

Answer 1

我认为没有办法使用极坐标表达式来生成组，因为它们总是依赖于前一个组。 话虽如此，这些组可以在 O(n) 中轻松生成，因此在 python 中这样做的惩罚应该很小。

import numpy as np
def make_groups(a, threshold=1.05):
    a=np.array(a)
    outarray=np.empty(len(a), dtype=a.dtype)
    outarray[0]=0
    curgroup=a[0]
    for indx, cur_a in enumerate(a[1:],1):
        if cur_a >= threshold * curgroup or cur_a * threshold <= curgroup:
            outarray[indx] = outarray[indx-1] + 1
            curgroup=cur_a
        else:
            outarray[indx] = outarray[indx-1]
    return pl.Series(outarray)

现在让我们将其应用到我们的数据中。

df = pl.DataFrame({'a': [100, 103, 105, 106, 105, 104, 103, 106, 100, 102]})

我们只是做了一个映射（顺便说一句，我尝试将 make_groups 制作成 np.ufunc 但无法让它工作）。

df \
    .with_columns(pl.col('a').map(lambda x: make_groups(x, 1.05)).alias('group')) \
    .with_columns((pl.col('a').list().over('group').arr.first()).alias('groupFirst'))

shape: (10, 3)
┌─────┬───────┬────────────┐
│ a   ┆ group ┆ groupFirst │
│ --- ┆ ---   ┆ ---        │
│ i64 ┆ i64   ┆ i64        │
╞═════╪═══════╪════════════╡
│ 100 ┆ 0     ┆ 100        │
│ 103 ┆ 0     ┆ 100        │
│ 105 ┆ 1     ┆ 105        │
│ 106 ┆ 1     ┆ 105        │
│ ... ┆ ...   ┆ ...        │
│ 103 ┆ 1     ┆ 105        │
│ 106 ┆ 1     ┆ 105        │
│ 100 ┆ 2     ┆ 100        │
│ 102 ┆ 2     ┆ 100        │
└─────┴───────┴────────────┘

顺便说一下，如果您只想使用默认阈值，那么您可以这样做...

df \
    .with_columns(pl.col('a').map(make_groups).alias('group')) \
    .with_columns((pl.col('a').list().over('group').arr.first()).alias('groupFirst'))

Answer 2

假设您想要在值超过值

6

时重置 cum_max。你可以这样做：

df = pl.DataFrame({'a': [1, 3, 5, 6, 1, 4, 3, 6, 5, 6]})

(
    df.with_columns(
        (pl.col("a") >= 6)
        .shift(1)
        .fill_null(False)
        .cum_sum()
        .alias("group")
    ).with_columns(
        pl.col("a")
        .cum_max()
        .over(pl.col("group"))
        .alias("cum_max")
    )
)

shape: (10, 3)
┌─────┬───────┬─────────┐
│ a   ┆ group ┆ cum_max │
│ --- ┆ ---   ┆ ---     │
│ i64 ┆ u32   ┆ i64     │
╞═════╪═══════╪═════════╡
│ 1   ┆ 0     ┆ 1       │
│ 3   ┆ 0     ┆ 3       │
│ 5   ┆ 0     ┆ 5       │
│ 6   ┆ 0     ┆ 6       │
│ 1   ┆ 1     ┆ 1       │
│ 4   ┆ 1     ┆ 4       │
│ 3   ┆ 1     ┆ 4       │
│ 6   ┆ 1     ┆ 6       │
│ 5   ┆ 2     ┆ 5       │
│ 6   ┆ 2     ┆ 6       │
└─────┴───────┴─────────┘