Python中基于4列滑动窗口高效更新DataFrame行的方法

我有一个仅包含0和1的DataFrame，希望通过宽度为4的列滑动窗口遍历所有行进行更新：当窗口内存在1时，将该窗口中包含1的行对应的整个窗口列全部设为1；若窗口中没有1则不做修改。目前通过循环实现该逻辑，但因DataFrame包含数百万行，效率极低，恳请提供高效解决方案。

示例DataFrame：

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, 1, 1]
df['a4'] = [0, 1, 1]
df['a5'] = [0, 0, 0]
df['a6'] = [0, 0, 1]
df['a7'] = [1, 0, 1]
df['a8'] = [1, 0, 1]
df['a9'] = [0, 0, 0]
df['a10'] = [0, 0, 0]
df['a11'] = [0, 0, 0]
df['a12'] = [0, 0, 0]

期望输出：

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

方案1：按不重叠列窗口处理（匹配示例输出） #

从示例输出判断，需求实际是将列划分为宽度为4的不重叠窗口（如a1-a4、a5-a8、a9-a12），当窗口内存在1时将整个窗口设为1。以下是向量化实现：

import pandas as pd

# 提取数值列（排除Ids）
num_cols = df.columns.drop('Ids')
# 将列分成每4个一组
col_groups = [num_cols[i:i+4] for i in range(0, len(num_cols), 4)]

# 遍历每个列组，若组内有1则将整组设为1
for group in col_groups:
    # 判断每行该组是否包含1
    has_one = df[group].any(axis=1)
    # 对有1的行，将该组所有列设为1
    df.loc[has_one, group] = 1

print(df)

该方法利用pandas原生向量化操作，避免逐行循环，处理百万级数据的效率远高于循环实现。

方案2：步长为1的滑动窗口（符合原描述） #

如果需求是步长为1的滑动窗口（窗口依次为a1-a4、a2-a5、a3-a6...），只要列被任何包含1的窗口覆盖就设为1，可通过numpy的滑动窗口视图实现：

import pandas as pd
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

# 提取数值列数据为numpy数组
num_data = df.drop('Ids', axis=1).values
window_size = 4

# 生成每行的滑动窗口（形状：(行数, 窗口数, 窗口大小)）
sliding_windows = sliding_window_view(num_data, window_size, axis=1)
# 判断每个窗口是否包含1（形状：(行数, 窗口数)）
window_has_one = (sliding_windows == 1).any(axis=2)

# 初始化结果数组为原数据副本
result = num_data.copy()
# 遍历每个窗口，将对应位置设为1
for i in range(window_has_one.shape[1]):
    cols = slice(i, i+window_size)
    result[window_has_one[:, i], cols] = 1

# 将结果合并回原DataFrame
df[num_cols] = result
print(df)

该方法借助numpy的底层优化，完全避免Python层面的循环，适合处理超大规模数据集。

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