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

Question

我有一个以下结构的DataFrame：

    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中出现1s时没有模式。

我现在的代码是：

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