python - 如何根据两列获取累计计数

Question

假设我们有以下数据框。如果我们想找到连续 1 的数量，您可以使用下面的代码。

    col
0   0
1   1
2   1
3   1
4   0
5   0
6   1
7   1
8   0
9   1
10  1
11  1
12  1
13  0
14  1
15  1


df['col'].groupby(df['col'].diff().ne(0).cumsum()).cumsum()

但我看到的问题是当您需要将 groupby 与 id 字段一起使用时。如果我们向数据帧添加一个 id 字段(如下)，则会变得更加复杂。我们不能再使用上面的解决方案。

    id  col
0   B   0
1   B   1
2   B   1
3   B   1
4   A   0
5   A   0
6   B   1
7   B   1
8   B   0
9   B   1
10  B   1
11  A   1
12  A   1
13  A   0
14  A   1
15  A   1

当遇到这个问题时，我看到了制作一个帮助器系列以在 groupby 中使用的案例，如下所示:

s = df['col'].eq(0).groupby(df['id']).cumsum()
df['col'].groupby([df['id'],s]).cumsum()

这有效，但问题是第一组包含第一行，这不符合条件。这通常不是问题，但如果我们想找到计数，那就是问题了。将最后一个 groupby() 末尾的 cumsum() 替换为 .transform('count') 实际上会得到 6 而不是 5 作为第一个 B 组中连续 1 的计数。

我能想到的解决这个问题的唯一解决方案是以下代码:

df['col'].groupby([df['id'],df.groupby('id')['col'].transform(lambda x: x.diff().ne(0).astype(int).cumsum())]).transform('count')

预期输出:

0     1
1     5
2     5
3     5
4     2
5     2
6     5
7     5
8     1
9     2
10    2
11    2
12    2
13    1
14    2
15    2

这可行，但使用 transform() 两次，我听说这不是最快的。这是我能想到的唯一使用 diff().ne(0) 获取“真实”组的解决方案。

索引1,2,3,6和7都是id B，在'col'列中具有相同的值，因此计数不会被重置，因此它们都会分开属于同一组。

可以在不使用多个 .transform() 的情况下完成此操作吗？

Answer 1

• 以下代码仅使用 1 个 .transform()，并依赖于对索引进行排序来获取正确的计数。
  • 保留原始索引，因此最终结果可以重新索引回原始顺序。
• 使用 cum_counts['cum_counts'] 获得准确的所需输出，无需其他列。

import pandas as pd

# test data as shown in OP
df = pd.DataFrame({'id': ['B', 'B', 'B', 'B', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'A', 'A', 'A', 'A', 'A'], 'col': [0, 1, 1, 1, 0, 0, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1]})

# reset the index, then set the index and sort
df = df.reset_index().set_index(['index', 'id']).sort_index(level=1) 

          col
index id     
4     A     0
5     A     0
11    A     1
12    A     1
13    A     0
14    A     1
15    A     1
0     B     0
1     B     1
2     B     1
3     B     1
6     B     1
7     B     1
8     B     0
9     B     1
10    B     1

# get the cumulative sum
g = df.col.ne(df.col.shift()).cumsum()

# use g to groupby and use only 1 transform to get the counts
cum_counts = df['col'].groupby(g).transform('count').reset_index(level=1, name='cum_counts').sort_index()

      id  cum_counts
index               
0      B           1
1      B           5
2      B           5
3      B           5
4      A           2
5      A           2
6      B           5
7      B           5
8      B           1
9      B           2
10     B           2
11     A           2
12     A           2
13     A           1
14     A           2
15     A           2