pandas数据新增一列计算连续次数

求教大神,想要计算一下一列中每一个数字连续出现的次数,请问有什么简单的方法吗?因为原始数据非常大,所以希望性能好一点,谢谢~

df1 = pd.DataFrame({'X': [1, 2, 2, 1, 1, 1, 2, 1, 1, 1]})
df1

希望生成以下的结果
image

最佳答案

使用Python原生的 itertools 中的函数可以完成.

from itertools import groupby

groups = groupby(df1['X'].values.tolist())

grouped_elements = [list(group) for key, group in groups]

会得到到如下结果:

[[1], [2, 2], [1, 1, 1], [2], [1, 1, 1]]

接下来是得到每个group的长度:

grouped_cnt_list = [[len(i)] * len(i) for i in grouped_elements]

结果如下:

[[1], [2, 2], [3, 3, 3], [1], [3, 3, 3]]

grouped_cnt_list 这个变量flatten之后并作为新的column添加到dataframe中:

from itertools import chain
df1['Y'] = list(chain.from_iterable(grouped_cnt_list))

参考官方文档:

itertools 的 groupby, https://docs.python.org/3/lib...
itertools 的 chain.from_iterable, https://docs.python.org/3/lib...