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)

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

我知道我可以按 id 分组，然后从各个组中过滤行，但我不知道如何根据组使过滤器动态化。我试过循环遍历这些组，但速度很慢（现在我有 19,000 个 ID，但它只会随着我继续项目而增长）。

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

Answer 1

这是另一个解决方案：groupby + 自定义函数

def get_data(x):
    idx = np.where(x['product']=='b')[0]
    if idx.size:
        return x[idx[0]:]

out = df.groupby('id').apply(get_data).droplevel(0).sort_index()

输出：

    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

Answer 2

您可以使用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']）