高效计算DataFrame列中各位置之后的连续NaN值数量

嘿，这个需求挺实用的——要算出DataFrame里每个位置开始往后连续出现的NaN个数对吧？我有个高效的实现方案，用pandas的内置函数就能搞定，完全不用写低效的逐行循环，一起来看看：

步骤1：准备测试数据 #

先把你给出的输入DataFrame构造出来，方便我们验证结果：

import pandas as pd
import numpy as np

df = pd.DataFrame({
    'A': [0.1880, 0.2510, np.nan, np.nan, np.nan, 0.2300, 0.1670, 0.0835, 0.0418, 0.0209, np.nan, np.nan, np.nan],
    'B': [0.345, 0.585, np.nan, np.nan, 1.150, 1.210, 1.290, 1.400, np.nan, np.nan, np.nan, np.nan, np.nan]
})

步骤2：定义处理函数 #

我们需要对每一列单独处理，从后往前统计连续NaN的数量，遇到非NaN值就重置计数：

def count_consecutive_nans_from_pos(col):
    # 标记当前列的NaN位置，NaN为True，非NaN为False
    is_nan = col.isna()
    # 反转布尔数组，从列的末尾开始处理
    reversed_is_nan = is_nan[::-1]
    # 创建分组键：每遇到一个非NaN值，就生成一个新的分组（用于重置计数）
    groups = (~reversed_is_nan).cumsum()
    # 每个分组内累加计数，得到反转后的连续NaN数量
    counts = reversed_is_nan.groupby(groups).cumsum()
    # 反转回原顺序，得到每个位置对应的连续NaN数量
    return counts[::-1]

步骤3：应用函数到整个DataFrame #

直接用apply方法把函数作用到每一列，就能得到目标结果：

result_df = df.apply(count_consecutive_nans_from_pos)
print(result_df)

运行后输出的结果就和你给出的示例完全一致：

A  B
0   0  0
1   0  0
2   3  2
3   2  1
4   1  0
5   0  0
6   0  0
7   0  0
8   0  5
9   0  4
10  3  3
11  2  2
12  1  1

原理简单解释 #

• 我们先把列反转，这样可以从末尾开始统计连续NaN，遇到非NaN就触发分组重置；
• 通过cumsum()生成分组键，每个连续的NaN块会被分到同一个组里；
• 组内的累加计数，反转后刚好对应原顺序中从当前位置到下一个非NaN值之间的连续NaN数量。

这个方法用了pandas的矢量化操作，比逐行循环快得多，即使处理大尺寸的DataFrame也能保持高效。

备注：内容来源于stack exchange，提问作者Vitamin C