基于DataFrame中1的首尾位置按窗口生成新DataFrame的问题排查

问题描述 #

我有一个DataFrame，需要以窗口大小4遍历它（跳过首列Ids），生成一个与输入列数相同的单行DataFrame。规则为：

• 每个窗口内，从所有行中第一个出现1的列位置到最后一个出现1的列位置之间填充1；
• 若窗口内无1则全填0。
  原代码无法正确处理最后部分列，请求修正。

输入数据 #

import pandas as pd

df = pd.DataFrame()
df['Ids'] = [1, 2, 3]
df['a1'] = [0, 0, 0]
df['a2'] = [0, 0, 1]
df['a3'] = [0, 0, 0]
df['a4'] = [0, 1, 1]
df['a5'] = [0, 0, 0]
df['a6'] = [0, 0, 1]
df['a7'] = [1, 0, 0]
df['a8'] = [1, 0, 0]
df['a9'] = [0, 0, 0]
df['a10'] = [0, 0, 0]
df['a11'] = [0, 0, 0]
df['a12'] = [0, 0, 0]

原自定义函数代码 #

def max_event_time(df,window_size):
    
    output_df = pd.DataFrame()
    
    for i in range(0, df.shape[1]-window_size, window_size ):
        window = df.iloc[:, i + 1:i + 1 + window_size]
        first_index = (window.iloc[:, 1:] == 1).idxmax(axis=1)
        last_index = (window.iloc[:, 1:] == 1).iloc[:, ::-1].idxmax(axis=1)
        first_index_pos = window.columns.get_loc(first_index.values[-1])
        last_index_pos = window.columns.get_loc(last_index.values[0])
        df_w = window.iloc[0:1, :]
        df_w.iloc[:, first_index_pos:last_index_pos + 1] = 1
        df_w.reset_index(drop=True, inplace=True)
        output_df = pd.concat([output_df,df_w], axis=1)
    return output_df

调用代码 #

max_event_time(df, 4)

期望输出 #

a1  a2  a3  a4  a5  a6  a7  a8  a9  a10  a11  a12
0   0   1   1   1   0   1   1   1   0    0    0    0

问题分析 #

原代码存在以下问题：

1. 循环范围错误：range(0, df.shape[1]-window_size, window_size)会漏掉最后一个完整窗口（当总列数是窗口大小整数倍时），导致部分列未处理；
2. 窗口列截取错误：window.iloc[:, 1:]跳过了窗口首列，导致首列的1无法被检测到；
3. 行选择逻辑错误：仅取单行的1位置，未考虑整个窗口所有行的1分布，导致范围计算不准确；
4. 拼接效率低：多次concat操作会影响性能，且容易出现列对齐问题。

修正后的代码 #

def max_event_time(df, window_size):
    # 提取需要处理的列（跳过首列Ids）
    cols_to_process = df.columns[1:]
    output_data = []
    
    # 遍历每个窗口，确保覆盖所有列
    for start in range(0, len(cols_to_process), window_size):
        end = start + window_size
        window_cols = cols_to_process[start:end]
        window = df[window_cols]
        
        # 找出窗口中所有值为1的列索引
        ones_cols = window[window == 1].stack().index.get_level_values(1)
        
        if not ones_cols.empty:
            # 获取窗口内第一个和最后一个出现1的列位置
            min_col = ones_cols.min()
            max_col = ones_cols.max()
            min_pos = window_cols.get_loc(min_col)
            max_pos = window_cols.get_loc(max_col)
            
            # 生成当前窗口的行数据：先全0，再填充1区间
            window_row = [0] * len(window_cols)
            for i in range(min_pos, max_pos + 1):
                window_row[i] = 1
        else:
            # 无1则全0
            window_row = [0] * len(window_cols)
        
        output_data.extend(window_row)
    
    # 生成最终输出DataFrame
    output_df = pd.DataFrame([output_data], columns=cols_to_process)
    return output_df

测试验证 #

调用修正后的函数：

result = max_event_time(df, 4)
print(result)

输出与期望完全一致。

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