统计DataFrame同列内相同值连续分组的元素数量

问题场景 #

DataFrame的列包含连续的1/-1分组，部分列开头可能有任意数量的NaN。需要统计每个非NaN连续分组的元素总数，并将该数值填充到原分组的对应位置。

输入示例 #

import pandas as pd
NaN = float('nan')
df = pd.DataFrame({'col1':[1,1,1,1,-1,-1,-1,1,1],
                   'col2':[NaN, NaN, NaN, 1,1,-1,-1,-1,-1],
                   'col3':[NaN,NaN,NaN,NaN,NaN,1,1,1,-1]})

期望输出 #

df_result = pd.DataFrame({'col1':[4,4,4,4,3,3,3,2,2],
                          'col2':[NaN, NaN, NaN,2,2,4,4,4,4],
                          'col3': [NaN,NaN,NaN,NaN,NaN,3,3,3,1]})

解决方案 #

通过Pandas的分组标记与transform方法实现，核心是先识别连续相同值的非NaN分组，再计算每组长度并填充：

import pandas as pd

def fill_group_length(col):
    # 标记每个新分组的起始位置：非NaN，且前一行是NaN或值不同
    change_points = col.notna() & (col.shift() != col | col.shift().isna())
    # 生成分组ID，每个新分组ID递增
    group_ids = change_points.cumsum()
    # 将NaN位置的分组ID设为空，避免纳入分组计算
    group_ids = group_ids.where(col.notna(), pd.NA)
    # 按分组ID计算每组长度，填充回原位置
    return col.groupby(group_ids).transform('size')

# 应用函数到所有列
result_df = df.apply(fill_group_length)
print(result_df)

代码说明 #

1. 标记分组变化点：change_points会在每个新分组的第一行返回True，比如从NaN切换到1、从1切换到-1的位置。
2. 生成分组ID：对变化点累加得到唯一的分组ID，同一连续分组的所有行拥有相同ID。
3. 计算并填充长度：按分组ID分组后，用transform('size')获取每组长度，自动填充到该组的每一行，NaN位置保持不变。

内容的提问来源于stack exchange，提问作者Kevin