Pandas - 创建一个新列，其值在旧列中每次出现值 X 时都会增加

Question

我有一个具有以下结构的数据框：

    A  

0   1   
1   2   
2   3   
3   1   
4   2   
5   1   
6   2
7   3

现在我想创建一个新的 B 列，从 0 开始，每次 A 列中出现 1 时，其值都会递增。因此上面的数据框应如下所示：

    A   B

0   1   0
1   2   0
2   3   0
3   1   1
4   2   1
5   1   2 
6   2   2
7   3   2

请注意，A 中 1 的出现没有模式。

我现在的代码是：

def _add_col_B(data):

  data['B'] = -1

  ones = list((data.index[data['A'] == 1]))
  ones.append(len(data))

  sent = 0
  for i in range(len(ones)-1):
      data.loc[ones[i] : ones[i+1],:]['B'] = sent
      sent = sent + 1

  return data

%timeit -r 3 _add_col_B(data)
10 loops, best of 3: 184 ms per loop

但在我看来，它非常慢，特别是考虑到我需要重复执行此操作并且对于非常大的数据帧。有没有矢量化的方法来做到这一点？

Answer 1

采用矢量化方法，您可以编写：

df['B'] = (df['A'] == 1).cumsum() - 1

这会产生 DataFrame：

   A  B
0  1  0
1  2  0
2  3  0
3  1  1
4  2  1
5  1  2
6  2  2
7  3  2