如何用向量化方法求Pandas DataFrame每行连续非零列的最大数量

高效向量化解决方案 #

针对400万行、12列的DataFrame，我们可以通过以下步骤快速计算每行中连续大于0的列的最大数量：

1. 生成布尔掩码：将DataFrame转换为布尔数组，标记出所有值大于0的位置
2. 计算连续计数：逐列维护每行的连续非零值计数，遇到0则重置为0
3. 取每行最大值：对每行的连续计数结果取最大值，得到目标列

具体代码实现：

import pandas as pd
import numpy as np

# 示例数据（替换为你的实际DataFrame即可）
df = pd.DataFrame(np.array([[284.77, 234.37, 243.8, 84.36, 0., 0., 0., 55.04, 228.2, 181.97, 0., 0.],
                            [13.78, 0., 38.58, 33.16, 0., 38.04, 74.02, 45.74, 27.2, 9.19, 0., 0.],
                            [88.66, 255.72, 323.19, 7.24, 0., 73.38, 45.73, 0., 0., 77.39, 26.57, 279.34],
                            [0., 0., 34.42, 9.16, 0., 43.4, 42.17, 123.69, 60.5, 25.47, 72.32, 7.29],
                            [320.6, 1445.56, 856.23, 371.21, 0., 244.22, 134.58, 631.59, 561.82, 1172.44, 895.68, 186.28],
                            [0., 0., 32.29, 1000.91, 0., 680., 585.46, 466.6, 0., 493.48, 157.1, 125.31]]),
                  columns=[1,2,3,4,5,6,7,8,9,10,11,12])

# 步骤1：生成布尔掩码数组
mask_np = (df > 0).to_numpy()

# 步骤2：计算每行连续非零值的计数
consec_counts = np.zeros_like(mask_np, dtype=np.int32)
consec_counts[:, 0] = mask_np[:, 0].astype(np.int32)

# 仅循环12次（对应列数），核心计算为numpy向量化操作
for col in range(1, mask_np.shape[1]):
    consec_counts[:, col] = np.where(mask_np[:, col], consec_counts[:, col-1] + 1, 0)

# 步骤3：添加结果列
df['MAX_CONSECUTIVE_COL'] = consec_counts.max(axis=1)

print(df['MAX_CONSECUTIVE_COL'])

方案优势 #

• 性能拉满：仅需循环12次，核心逻辑依赖numpy向量化运算，处理400万行数据仅需数秒
• 内存高效：使用numpy数组存储中间结果，比纯pandas操作节省内存开销
• 逻辑直观：逐列维护连续计数，代码简洁易懂，便于调试和维护

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