DataFrame binary列向量化优化报错求助：替换前两行0为1

问题背景 #

现有如下示例DataFrame：

df = pd.DataFrame({
    'time': pd.date_range(start='2023-03-15 00:00:00', periods=15, freq='1min'),
    'binary': [0,0,0,0,1,0,0,0,0,1,1,1,1,1,0],
    'nonbinary': np.random.randint(1, 10, 15)
})

需求是：找到binary列中的所有1（单个或连续出现），将每个1对应的前两分钟（即前两行）的0改为1。变换效果如下：

原数据: [0,0,0,0,1,0,0,0,0,1,1,1,1,1,0]
目标结果: [0,0,1,1,1,0,0,0,1,1,1,1,1,1,1]

已编写的循环代码可正常运行并得到预期结果，但面对数千行的实际DataFrame时效率较低：

import pandas as pd
import numpy as np

# 创建DataFrame
df = pd.DataFrame({
    'time': pd.date_range(start='2023-03-15 00:00:00', periods=15, freq='1min'),
    'binary': [0,0,1,0,1,0,0,0,0,1,1,1,1,1,0],
    'nonbinary': np.random.randint(1, 10, 15)
})

# 复制binary列用于存储更新后的值
updated_binary = df['binary'].copy()

# 遍历每一行
for i in range(len(df)):
    # 检查当前行是否为1
    if df.loc[i, 'binary'] == 1:
        # 将当前行及前两行设为1
        updated_binary[max(0, i-2):i+1] = 1

# 更新原列
df['binary'] = updated_binary
print(df)

尝试向量化实现时出现报错：

import pandas as pd
import numpy as np

# 创建DataFrame
df = pd.DataFrame({
    'time': pd.date_range(start='2023-03-15 00:00:00', periods=15, freq='1min'),
    'binary': [0,0,0,0,1,0,0,0,0,1,1,1,1,1,0],
    'nonbinary': np.random.randint(1, 10, 15)
})

updated_binary = df['binary'].copy()
Ones_positions = df['binary'].eq(1).values

window = np.zeros(3, dtype=bool)
window[:2] = True

zero_positions = np.logical_and.reduce((
    np.roll(Ones_positions, 2), 
    np.logical_not(Ones_positions),
    np.convolve(window, np.ones(3, dtype=bool), mode='valid') == 2
), axis=0)

updated_binary[zero_positions] = 1
df['binary'] = updated_binary
print(df)

报错信息：

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

错误原因 #

报错的核心是np.convolve使用mode='valid'返回的数组长度（13）与另外两个数组（长度15）不一致，导致np.logical_and.reduce无法对维度不匹配的数组做逻辑运算。此外原思路过于复杂，可通过更简洁的向量化方法实现需求。

高效向量化解决方案 #

方法1：Numpy卷积实现 #

利用卷积核标记所有需要设为1的位置，完全向量化处理，效率极高：

import pandas as pd
import numpy as np

# 创建示例DataFrame
df = pd.DataFrame({
    'time': pd.date_range(start='2023-03-15 00:00:00', periods=15, freq='1min'),
    'binary': [0,0,0,0,1,0,0,0,0,1,1,1,1,1,0],
    'nonbinary': np.random.randint(1, 10, 15)
})

# 定义卷积核：每个1会激活自身及前两个位置
kernel = np.ones(3, dtype=int)
# 执行卷积，mode='same'保证输出长度与原数组一致
activated = np.convolve(df['binary'].values, kernel, mode='same') >= 1
# 将布尔数组转为整数，得到目标binary列
df['binary'] = activated.astype(int)

print(df)

方法2：Pandas Rolling窗口实现 #

通过反向滚动窗口标记后续是否存在1，实现需求：

import pandas as pd
import numpy as np

# 创建示例DataFrame
df = pd.DataFrame({
    'time': pd.date_range(start='2023-03-15 00:00:00', periods=15, freq='1min'),
    'binary': [0,0,0,0,1,0,0,0,0,1,1,1,1,1,0],
    'nonbinary': np.random.randint(1, 10, 15)
})

# 反转列后用rolling窗口取最大值，判断后续2个位置内是否有1
has_future_one = df['binary'][::-1].rolling(window=3, min_periods=1).max()[::-1] == 1
# 转换为整数列
df['binary'] = has_future_one.astype(int)

print(df)

两种方法均无需循环，处理大型数据集时性能远优于逐行遍历。

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