转动时重置的累计数negative/positive

Cumulative sum that resets when turning negative/positive

enter image description here我正在尝试向我的 polars 数据框添加一列(C 列),该列计算数据框的其中一列(A 列)的值比 greater/less 的次数另一列(B 列)的值。一旦值从 less/greater 变为 greater/less,累积总和应重置并再次从 1/-1 开始计数。

不是很优雅或 Pythonic,但像下面这样的东西应该可以工作:

import pandas as pd

df = pd.DataFrame({'a': [10, 10, 10, 8, 8, 8, 15, 15]
,'b': [9, 9, 9, 9, 10, 10, 11, 11]})

df['c'] = df.apply(lambda row: 1 if row['a'] > row['b'] else 0, axis=1)
df['d'] = df.apply(lambda row: 0 if row['a'] > row['b'] else -1, axis=1)
for i in range(1, len(df)):
    if df.loc[i, 'a'] > df.loc[i, 'b']: 
        df.loc[i, 'c'] = df.loc[i-1, 'c']  + 1
        df.loc[i, 'd'] = 0
    else:
        df.loc[i, 'd'] = df.loc[i-1, 'd']  - 1
        df.loc[i, 'c'] = 0
        
df['ans'] = df['c'] + df['d']
print(df)

此外,您可能需要考虑当 a 列和 b 列相等时的特定情况下该值应该是多少。

数据

我将更改您提供的示例中的数据。

df = pl.DataFrame(
    {
        "a": [11, 10, 10, 10, 9, 8, 8, 8, 8, 8, 15, 15, 15],
        "b": [11, 9, 9, 9, 9, 9, 10, 8, 8, 10, 11, 11, 15],
    }
)
print(df)
shape: (13, 2)
┌─────┬─────┐
│ a   ┆ b   │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 11  ┆ 11  │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 10  ┆ 9   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 10  ┆ 9   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 10  ┆ 9   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 9   ┆ 9   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 8   ┆ 9   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 8   ┆ 10  │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 8   ┆ 8   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 8   ┆ 8   │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 8   ┆ 10  │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 15  ┆ 11  │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 15  ┆ 11  │
├╌╌╌╌╌┼╌╌╌╌╌┤
│ 15  ┆ 15  │
└─────┴─────┘

注意两列相同的情况。您的 post 没有说明在这些情况下该怎么做,所以我对应该发生的事情做了一些假设。 (您可以调整代码以不同方式处理这些情况。)

算法

df = (
    df
    .with_column((pl.col("a") - pl.col("b")).sign().alias("sign_a_minus_b"))
    .with_column(
        pl.when(pl.col("sign_a_minus_b") == 0)
        .then(None)
        .otherwise(pl.col("sign_a_minus_b"))
        .forward_fill()
        .alias("run_type")
    )
    .with_column(
        (pl.col("run_type") != pl.col("run_type").shift_and_fill(1, 0))
        .cumsum()
        .alias("run_id")
    )
    .with_column(pl.col("sign_a_minus_b").cumsum().over("run_id").alias("result"))
)
print(df)
shape: (13, 6)
┌─────┬─────┬────────────────┬──────────┬────────┬────────┐
│ a   ┆ b   ┆ sign_a_minus_b ┆ run_type ┆ run_id ┆ result │
│ --- ┆ --- ┆ ---            ┆ ---      ┆ ---    ┆ ---    │
│ i64 ┆ i64 ┆ i64            ┆ i64      ┆ u32    ┆ i64    │
╞═════╪═════╪════════════════╪══════════╪════════╪════════╡
│ 11  ┆ 11  ┆ 0              ┆ null     ┆ 1      ┆ 0      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 10  ┆ 9   ┆ 1              ┆ 1        ┆ 2      ┆ 1      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 10  ┆ 9   ┆ 1              ┆ 1        ┆ 2      ┆ 2      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 10  ┆ 9   ┆ 1              ┆ 1        ┆ 2      ┆ 3      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 9   ┆ 9   ┆ 0              ┆ 1        ┆ 2      ┆ 3      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8   ┆ 9   ┆ -1             ┆ -1       ┆ 3      ┆ -1     │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8   ┆ 10  ┆ -1             ┆ -1       ┆ 3      ┆ -2     │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8   ┆ 8   ┆ 0              ┆ -1       ┆ 3      ┆ -2     │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8   ┆ 8   ┆ 0              ┆ -1       ┆ 3      ┆ -2     │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 8   ┆ 10  ┆ -1             ┆ -1       ┆ 3      ┆ -3     │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 15  ┆ 11  ┆ 1              ┆ 1        ┆ 4      ┆ 1      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 15  ┆ 11  ┆ 1              ┆ 1        ┆ 4      ┆ 2      │
├╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
│ 15  ┆ 15  ┆ 0              ┆ 1        ┆ 4      ┆ 2      │
└─────┴─────┴────────────────┴──────────┴────────┴────────┘

我在输出中保留了中间计算,只是为了展示算法的工作原理。 (你可以放下它们。)

基本思路是为每个 运行 的正值或负值计算一个 run_id。然后,我们将使用 cumsum 函数和 over 窗口表达式在每个 run_id.

上创建 positives/negatives 的 运行ning 计数

关键假设ab 列中的关系不会中断 运行,但它们对总数没有贡献对于 positive/negative 个值中的 运行 个值。

sign_a_minus_b 做两件事:它识别 运行 是否是 positive/negative,以及 ab 列是否有平局。

run_type 扩展任何 运行 包括在列 ab 中出现平局的任何情况。列顶部的 null 值是预期的 - 它显示了当第一行出现平局时会发生什么。

result 是输出列。请注意,绑定的列不会中断 运行,但它们不会对 运行.

的总数做出贡献

最后一点:如果不允许 ab 列中的关系,则可以简化此算法...并且 运行 更快。