基于连续0值计数条件的DataFrame新列生成及代码优化需求

数据结构与字段说明 #

原始DataFrame每行对应1秒数据，结构如下：

Answers,all_answers,Score
1.0,1,1
0.0,0,1
0.0,0,2
0.0,0,3
-1.0,1,1
0.0,0,1
0.0,0,2
1.0,1,1
-1.0,1,2
0.0,0,1
1.0,1,1

各字段含义：

• Answers：对应秒数的计算符号，1为正，-1为负；
• all_answers：Answers的绝对值，用于生成Score列；
• Score：连续出现0或1的次数计数。

核心需求 #

生成New列，规则如下：

• 所有Answers为0的行，New直接设为0；
• 对于Answers为1/-1的行，仅当该值与前一个非0值之间的连续0的Score计数≤5，且与后一个非0值之间的连续0的Score计数≤5时，保留原Answers值；否则设为0。
  （注：首尾的非0值只需判断单侧的连续0计数）

示例场景展示 #

场景1：中间连续0超过5，两端非0值设为0 #

Answers,all_answers,Score,New
1.0,1,1,1
0.0,0,1,0
0.0,0,2,0
0.0,0,3,0
0.0,0,4,0
0.0,0,5,0
0.0,0,6,0
-1.0,1,1,0
0.0,0,1,0
0.0,0,2,0
1.0,1,1,1
-1.0,1,2,-1
0.0,0,1,0
1.0,1,1,1

场景2：连续0≤5，所有非0值保留 #

Answers,all_answers,Score,New
1.0,1,1,1
0.0,0,1,0
-1.0,1,1,-1
0.0,0,1,0
1.0,1,1,1
0.0,0,1,0
-1.0,1,1,-1
0.0,0,1,0
1.0,1,1,1
0.0,0,1,0
-1.0,1,1,-1
0.0,0,1,0
0.0,0,2,0
1.0,1,1,1

场景3：首尾长连续0后的非0值设为0 #

Answers,all_answers,Score,New
0.0,0,19,0
0.0,0,20,0
0.0,0,21,0
0.0,0,22,0
0.0,0,23,0
0.0,0,24,0
0.0,0,25,0
0.0,0,26,0
-1.0,1,1,0
0.0,0,1,0
0.0,0,2,0
0.0,0,3,0
0.0,0,4,0
0.0,0,5,0
0.0,0,6,0
1.0,1,1,0

场景4：单侧连续0≤5的非0值保留 #

Answers,all_answers,Score,New
1.0,1,1,0
0.0,0,1,0
0.0,0,2,0
0.0,0,3,0
0.0,0,4,0
0.0,0,5,0
-1.0,1,1,-1
1.0,1,2,1
-1.0,1,3,-1
1.0,1,4,1
-1.0,1,5,-1
0.0,0,1,0
0.0,0,2,0
0.0,0,3,0
0.0,0,4,0
1.0,1,1,1
0.0,0,1,0
0.0,0,2,0
0.0,0,3,0
0.0,0,4,0
0.0,0,5,0
0.0,0,6,0
0.0,0,7,0

现有代码问题 #

原代码逻辑错误，仅判断连续值的长度是否≥5，未针对两个非0值之间的连续0计数做判断，无法覆盖需求场景：

import numpy as np
import pandas as pd

s = np.sign(df['all_answers'])
group = s.ne(s.shift()).cumsum()
df['Score'] = s.groupby(group).cumcount().add(1)
g = df['all_answers'].ne(df['all_answers'].shift()).cumsum()
X1 = int(x)
m = df.groupby(g)['all_answers'].transform('size').ge(5)
df['New'] = df['Answers'].where(m, 0)

优化后的代码实现 #

import pandas as pd

def generate_new_column(df):
    # 复制原数据避免修改原表
    df = df.copy()
    # 标记非0行
    df['non_zero'] = df['Answers'].ne(0)
    # 获取所有非0行的索引
    non_zero_indices = df[df['non_zero']].index.tolist()
    
    # 初始化New列为0
    df['New'] = 0
    
    for i in range(len(non_zero_indices)):
        current_idx = non_zero_indices[i]
        # 计算与前一个非0值之间的连续0数量
        if i > 0:
            prev_idx = non_zero_indices[i-1]
            # 中间的0的最大Score值就是连续0的长度
            zero_count_between = df.loc[prev_idx+1:current_idx-1, 'Score'].max() if prev_idx+1 <= current_idx-1 else 0
        else:
            # 第一个非0值，前面没有非0，只需判断后面
            zero_count_between = 0
        
        # 计算与后一个非0值之间的连续0数量
        if i < len(non_zero_indices)-1:
            next_idx = non_zero_indices[i+1]
            zero_count_next = df.loc[current_idx+1:next_idx-1, 'Score'].max() if current_idx+1 <= next_idx-1 else 0
        else:
            # 最后一个非0值，后面没有非0，只需判断前面
            zero_count_next = 0
        
        # 判断是否满足条件：前后连续0计数都≤5（单侧的话只需存在的一侧≤5）
        if (i == 0 and zero_count_next <=5) or \
           (i == len(non_zero_indices)-1 and zero_count_between <=5) or \
           (zero_count_between <=5 and zero_count_next <=5):
            df.loc[current_idx, 'New'] = df.loc[current_idx, 'Answers']
    
    # 移除临时列
    df.drop('non_zero', axis=1, inplace=True)
    return df

# 调用示例
# df = pd.read_csv('your_data.csv')
# df = generate_new_column(df)

代码说明 #

1. 标记非0行：通过non_zero列筛选出所有非0值的位置；
2. 遍历非0行：逐个检查每个非0值与前后最近非0值之间的连续0的最大Score（即连续0的长度）；
3. 条件判断：根据首尾位置的特殊情况，判断连续0长度是否≤5，满足则保留原Answers值，否则设为0；
4. 性能优化：针对5000+行的DataFrame，该逻辑仅遍历非0行，效率较高（连续0平均300的情况下，非0行数量约为5000/300≈17行，遍历成本极低）。

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