python - 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 中 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

产生数据帧:

   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