取连续行直到特定值熊猫数据框

Question

我有一个类似于下面的熊猫数据框。 它包含 3 个标识符，分为相应的月份（长度可以变化）、它们相应的值和布尔标志。 对于每个标识符，我需要取值直到看到第一个“1”（包括），以及其他变量。 如果所有标志都是“0”，那么它将获取该 ID 的所有行。

所需的输出如下图所示。

    data = {'id':['123', '123', '123', '123', '123', '456', '456', '456', '456', '789', '789', '789', '789', '789', '789'],
            'month':[1,2,3,4,5,1,2,3,4,1,2,3,4,5,6],
            'value_1': [232,432,556,223,643,556,121,853,343,324,654,765,128,543,776],
            'flag':[0,0,0,1,1,0,0,0,0,0,0,0,0,1,1]}
     
    # Create DataFrame
    d = pd.DataFrame(data)

我曾尝试使用 groupby 进行转换（如下所示）。 对于值列，我只关心第一个值。 但是，我希望所有月份都按相同的顺序排列，这似乎用这种方法是不可能的。

temp = d['flag'].ne(1).cumsum()
grouped = d.groupby(temp).agg({'id': 'first',
                            'value_1': 'first',
                            'flag': lambda x: max(x)})

Answer 1

IIUC，尝试：

output = d[d.groupby("id")["flag"].transform(lambda x: x.shift().fillna(0).cumsum())==0]
output["value_1"] = output.groupby("id")["value_1"].transform("first")

>>> output

     id  month  value_1  flag
0   123      1      232     0
1   123      2      232     0
2   123      3      232     0
3   123      4      232     1
5   456      1      556     0
6   456      2      556     0
7   456      3      556     0
8   456      4      556     0
9   789      1      324     0
10  789      2      324     0
11  789      3      324     0
12  789      4      324     0
13  789      5      324     1