Python Groupby 为每个组单独筛选

Question

假设我有以下数据框：

import pandas as pd
d = {'id':  [1, 2, 3, 3, 3, 2, 2, 1, 2, 3, 2, 3],
'date': [1, 1, 1, 2, 2, 2, 3, 3, 3, 3, 4, 4],
'product': ['a', 'a', 'b', 'a', 'b', 'a', 'b', 'c', 'b', 'c', 'c', 'c']}
df = pd.DataFrame(d)

我想在他们购买产品"b"的当天和之后保留每个ID的所有数据，并在他们购买产品"b"之前删除所有数据。ID 1 没有数据，因为他们没有购买产品，ID 2 将包含第 3 天和第 4 天的数据，ID 3 将包含第 1-4 天的数据。

我知道我可以按 id 分组，然后过滤来自各个组的行，但我无法弄清楚如何根据组使过滤器动态。我尝试遍历组，但它很慢(现在我有 19,000 个 ID，但它只会随着我继续项目而增长)。

任何帮助将不胜感激。谢谢！

Answer 1

您可以选择带有eq的产品"b"，并使用groupby+cummax将每组的连续行设置为 True。然后对数据帧进行切片

df[df['product'].eq('b').groupby(df['id']).cummax()]

输出：

id  date product
2    3     1       b
3    3     2       a
4    3     2       b
6    2     3       b
8    2     3       b
9    3     3       c
10   2     4       c
11   3     4       c

注意。这假定数据帧是按日期排序的。如果不使用sort_values(by='date')(或by=['group', 'date'])