如何在Polars或Pandas中按条件实现累计求和？

Polars 实现方案 #

思路 #

核心是将每个报告日（is_reporting_day?为True）及其之前连续的非报告日划分为同一分组，对每组的customers求和后，仅保留报告日的行，最后计算累计总和。具体步骤：

1. 生成分组标识：反向累计报告日的True值后再反转，让每个报告日和它之前的非报告日归为同一组。
2. 按分组标识对customers求和，得到每组的总客户数。
3. 筛选出原数据中is_reporting_day?为True的行，合并分组求和结果。
4. 对求和后的结果计算累计总和，生成cum_sum列。

代码实现 #

import polars as pl

df = pl.from_repr("""
┌─────────────────────┬───────────┬───────────────────┐
│ date                ┆ customers ┆ is_reporting_day? │
│ ---                 ┆ ---       ┆ ---               │
│ datetime[ns]        ┆ i64       ┆ bool              │
╞═════════════════════╪═══════════╪═══════════════════╡
│ 2022-01-01 00:00:00 ┆ 3         ┆ true              │
│ 2022-01-02 00:00:00 ┆ 4         ┆ false             │
│ 2022-01-03 00:00:00 ┆ 5         ┆ false             │
│ 2022-01-04 00:00:00 ┆ 3         ┆ false             │
│ 2022-01-05 00:00:00 ┆ 2         ┆ true              │
└─────────────────────┴───────────┴───────────────────┘
""")

# 生成分组键：每个报告日及之前的非报告日为一组
group_key = pl.col("is_reporting_day?").reverse().cumsum().reverse()

# 分组求和，合并回原数据，筛选报告日并计算累计和
result = (
    df
    .with_columns(group_key=group_key)
    .with_columns(
        group_sum=pl.col("customers").sum().over("group_key")
    )
    .filter(pl.col("is_reporting_day?"))
    .with_columns(cum_sum=pl.col("group_sum").cumsum())
    .select(["date", "customers", "is_reporting_day?", "cum_sum"])
)

print(result)

输出结果 #

┌─────────────────────┬───────────┬───────────────────┬─────────┐
│ date                ┆ customers ┆ is_reporting_day? ┆ cum_sum │
│ ---                 ┆ ---       ┆ ---               ┆ ---     │
│ datetime[ns]        ┆ i64       ┆ bool              ┆ i64     │
╞═════════════════════╪═══════════╪═══════════════════╪═════════╡
│ 2022-01-01 00:00:00 ┆ 3         ┆ true              ┆ 3       │
│ 2022-01-05 00:00:00 ┆ 2         ┆ true              ┆ 14      │
└─────────────────────┴───────────┴───────────────────┴─────────┘

Pandas 实现方案 #

思路 #

和Polars逻辑一致，通过反向累计报告日生成分组键，分组求和后筛选报告日，最后计算累计总和。

代码实现 #

import pandas as pd
import polars as pl

df = pl.from_repr("""
┌─────────────────────┬───────────┬───────────────────┐
│ date                ┆ customers ┆ is_reporting_day? │
│ ---                 ┆ ---       ┆ ---               │
│ datetime[ns]        ┆ i64       ┆ bool              │
╞═════════════════════╪═══════════╪═══════════════════╡
│ 2022-01-01 00:00:00 ┆ 3         ┆ true              │
│ 2022-01-02 00:00:00 ┆ 4         ┆ false             │
│ 2022-01-03 00:00:00 ┆ 5         ┆ false             │
│ 2022-01-04 00:00:00 ┆ 3         ┆ false             │
│ 2022-01-05 00:00:00 ┆ 2         ┆ true              │
└─────────────────────┴───────────┴───────────────────┘
""").to_pandas()

# 生成分组键
df["group_key"] = df["is_reporting_day?"][::-1].cumsum()[::-1]

# 分组求和
df["group_sum"] = df.groupby("group_key")["customers"].transform("sum")

# 筛选报告日，计算累计和
result = (
    df[df["is_reporting_day?"]]
    .assign(cum_sum=lambda x: x["group_sum"].cumsum())
    .loc[:, ["date", "customers", "is_reporting_day?", "cum_sum"]]
)

print(result)

输出结果 #

date  customers  is_reporting_day?  cum_sum
0 2022-01-01 00:00:00          3               True        3
4 2022-01-05 00:00:00          2               True       14

内容的提问来源于stack exchange，提问作者nam0_0