基于仅含0/1的DataFrame生成保留连续1前两位的同结构DataFrame

需求说明 #

给定一个仅包含0和1的Pandas DataFrame，生成结构完全相同的新DataFrame。要求对每一行中连续出现的1，仅保留前两个，后续的连续1全部转为0。例如输入序列[0,0,1,1,1,1,0,0]，转换后为[0,0,1,1,0,0,0,0]。

示例输入 #

import pandas as pd

df_dr = pd.DataFrame()
df_dr['0'] = [0,0]
df_dr['1'] = [1,0]
df_dr['2'] = [1,1]
df_dr['3'] = [1,1]
df_dr['4'] = [0,1]
df_dr['5'] = [1,0]
df_dr['6'] = [1,0]
df_dr['7'] = [1,0]
df_dr['8'] = [1,1]
df_dr['9'] = [0,1]
df_dr['10'] = [0,1]
df_dr['11'] = [0,1]
df_dr['12'] = [0,1]

预期输出 #

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

解决方案代码 #

import pandas as pd

def process_continuous_ones(df):
    # 定义单行处理逻辑
    def process_row(row):
        # 为连续相同元素生成分组ID
        groups = (row != row.shift()).cumsum()
        # 统计每个分组内的元素位置序号
        counts = groups.groupby(groups).cumcount() + 1
        # 保留0或分组内前2个1，其余转为0
        return row.where((row == 0) | (counts <= 2), 0)
    
    # 对每一行应用处理逻辑
    return df.apply(process_row, axis=1)

# 处理示例数据并打印结果
result_df = process_continuous_ones(df_dr)
print(result_df)

代码逻辑解释 #

1. 分组连续序列：通过row != row.shift()判断当前元素与前一个元素是否不同，再用cumsum()生成连续相同元素的分组ID，每一段连续的0或1都会被分到独立分组。
2. 组内位置计数：利用groupby(groups).cumcount()对每个分组内的元素从0开始计数，加1后得到元素在分组内的位置序号。
3. 过滤保留目标元素：使用where方法，仅保留值为0的元素，或是分组内前2个1的元素，其余元素统一转为0。

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