Pandas 条件 groupby：使用谓词对分区进行聚合

Question

我已经多次执行此任务，因此我想知道是否可以在使用谓词时以 OVER PARTITION BY 样式进行聚合。

关于窗口函数

groupby

的 SO 有数十个示例，并且不少使用条件，但没有一个条件取决于列的值（即不基于文字 - 常量值）。

考虑这个示例数据框：

import pandas as pd

df = pd.DataFrame(data={
    'id': ['0001', '0001', '0001', '0002', '0002', '0003', '0003', '0004'],
    'days': [2, 5, 6, 3, 5, 4, 6, 8],
    'amount': [100, 150, 150, 200, 100, 300, 250, 200]
})

print(df)

     id  days  amount
0  0001     2     100
1  0001     5     150
2  0001     6     150
3  0002     3     200
4  0002     5     100
5  0003     4     300
6  0003     6     250
7  0004     8     200

我想要按

amount

上的条件对

days

的窗口总和进行分组。此条件在

days

上可能不太相等，即给定多个

days

，计算所有

amount

的

days

之和，直到该行中

days

的值，即

df['days'] <= i

为每行

i

.

除了使用循环之外，我找不到任何其他解决方案：

df['run_sum_le'] = 0
for i in df.sort_values('days', ascending=False)['days'].unique():
    le = df['days'] <= i
    df.loc[le, 'run_sum_le'] = df[le]['amount'].sum()

print(df)

     id  days  amount  run_sum_le
0  0001     2     100         100
1  0001     5     150         850
2  0001     6     150        1250
3  0002     3     200         300
4  0002     5     100         850
5  0003     4     300         600
6  0003     6     250        1250
7  0004     8     200        1450

请注意，迭代器的正确排序是必要的。

另一个更清晰的例子：让我们使用大于等于和等于的条件来计算

amount

的总和，即对每行

amount

使用

df['days'] >= i

和

df['days'] == i

求和

i

。

df['run_sum_eq'] = 0
df['run_sum_ge'] = 0
for i in df.sort_values('days', ascending=True)['days'].unique():
    eq = df['days'] == i
    ge = df['days'] >= i
    df.loc[eq, 'run_sum_eq'] = df[eq]['amount'].sum()
    df.loc[ge, 'run_sum_ge'] = df[ge]['amount'].sum()

print(df)

     id  days  amount  run_sum_le  run_sum_eq  run_sum_ge
0  0001     2     100         100         100        1450
1  0001     5     150         850         250         850
2  0001     6     150        1250         400         600
3  0002     3     200         300         200        1350
4  0002     5     100         850         250         850
5  0003     4     300         600         300        1150
6  0003     6     250        1250         400         600
7  0004     8     200        1450         200         200

注 1：更改了迭代器中的顺序。 注 2：这种

eq = df['days'] == i

的方法是一种矫枉过正，因为可以用一个衬垫代替：

df['run_sum_eq'] = df['days'].map(df.groupby('days')['amount'].sum())

这只是通过

days

进行摸索，即它对应于常规 PARTITION BY 计算。

最后一个方法相当于：

df['days'].map(df.groupby('days')['amount'].agg(pd.Series.sum))

这让我想到我可以在

agg()

中使用 lambda，其中我可以包含谓词：

df['days'].map(df.groupby('days')['amount'].agg([('amount', lambda val: ???)]))

这就是我陷入困境的地方。

Answer 1

您希望它有多通用？

解决您的具体案例非常简单......

• LE只是一个cumsum
• EQ只是一个总和
• GE 是

  total - LE + EQ

这给出了...

le = (
  df
    .sort_values('days')
    .groupby('days')['amount']
    .sum()
    .cumsum()
    .rename('le')
)

df = (
  df
    .merge(le, on='days')
    .sort_index()
)

df['eq'] = (
  df
    .sort_values('days')
    .groupby('days')['amount']
    .transform(sum)
)

df['ge'] = df['amount'].sum() - df['le'] + df['eq']

https://trinket.io/python3/9aa879ac91