按组统计二进制列连续1的计数，求大数据高效处理方案

需求规则 #

• 按city_id分组统计binary列中连续值为1的计数
• 核心规则：
  1. 当binary值为0时重置计数器
  2. 切换至新分组时重置计数器
  3. 首次出现binary=1时计数从1开始

数据示例 #

data = {'city_id': [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6],
        'week': [1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7, 1, 2, 3, 4, 5, 6, 7],
        'binary': [0, 1, 1, 1, 0, 0, 0, 0, 1, 1, 0, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 1, 0, 1]}

df = pd.DataFrame(data)

原低效循环方案 #

df['consecutive'] = 0
for city in df['city_id'].unique():
    city_df = df[df['city_id'] == city]
    consecutive_count = 0
    for i in range(len(city_df)):
        if city_df['binary'].iloc[i] == 1:
            consecutive_count += 1
        else:
            consecutive_count = 0
        df.loc[(df['city_id'] == city) & (df['week'] == city_df['week'].iloc[i]), 'consecutive'] = consecutive_count

问题 #

该方案处理约250万条记录的大数据时效率极低，时常超时或运行数小时，急需更高效的实现方法。

高效实现方案 #

利用Pandas矢量化分组操作，完全避免Python循环，处理百万级数据仅需几秒到几十秒：

# 1. 按city_id分组，标记每个连续1的区块（binary为0时区块号递增）
df['block'] = df.groupby('city_id')['binary'].apply(lambda x: (x == 0).cumsum())

# 2. 按city_id+block分组，对每个区块内的行累加计数
df['consecutive'] = df.groupby(['city_id', 'block']).cumcount() + 1

# 3. 将binary为0的位置重置为0
df.loc[df['binary'] == 0, 'consecutive'] = 0

# 可选：删除临时区块列
df = df.drop('block', axis=1)

验证结果 #

处理示例数据后，consecutive列的核心结果片段：

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