如何高效获取二维Numpy数组每行首段连续NaN的最后索引？

需求说明 #

处理二维Numpy数组，每行按以下规则处理NaN：

• 若该行不以NaN开头，结果为 -1；
• 若该行以NaN开头，结果为该行起始连续未中断NaN序列的最后一个NaN的索引。

需处理数百万行的数组，因此需要最优实现方案。

示例数组 #

import numpy as np

arr = np.array([[1,11,np.nan,111,1111],
               [np.nan, np.nan, np.nan, 2, 22],
               [np.nan, np.nan, 3, 33, np.nan],
               [4, np.nan, np.nan, 44, 444],
               [np.nan, 5, 55, np.nan, 555],
               [np.nan, np.nan, np.nan, np.nan, np.nan]])

预期结果 #

result = [-1, 2, 1, -1, 0, 4]

现有实现代码 #

result = []

for i in range(arr.shape[0]):
    if np.isnan(arr[i])[0] == False:
        result += [-1]
    elif np.all(np.isnan(arr[i])):
        result += [arr.shape[1]-1]
    else:
        result += [np.where(np.isnan(arr[i]) == False)[0][0] - 1]

优化实现方案 #

现有代码采用逐行循环处理，在百万行数据场景下效率较低。以下是完全向量化的实现，利用Numpy的C级操作大幅提升处理速度：

import numpy as np

arr = np.array([[1,11,np.nan,111,1111],
               [np.nan, np.nan, np.nan, 2, 22],
               [np.nan, np.nan, 3, 33, np.nan],
               [4, np.nan, np.nan, 44, 444],
               [np.nan, 5, 55, np.nan, 555],
               [np.nan, np.nan, np.nan, np.nan, np.nan]])

# 生成每行的NaN掩码
nan_mask = np.isnan(arr)
# 判断每行是否以NaN开头
starts_with_nan = nan_mask[:, 0]
# 找到每行第一个非NaN的索引，全NaN时返回0
first_non_nan = np.argmax(~nan_mask, axis=1)
# 处理全NaN的情况：将全NaN行的first_non_nan设为列数，方便后续计算
all_nan = np.all(nan_mask, axis=1)
first_non_nan[all_nan] = arr.shape[1]

# 按规则生成结果
result = np.where(starts_with_nan, first_non_nan - 1, -1)
print(result.tolist())  # 输出 [-1, 2, 1, -1, 0, 4]

优化说明 #

• 避免了Python层的逐行循环，全部操作基于Numpy的向量化接口，利用底层C实现加速，处理百万级数据时性能提升显著；
• 通过argmax快速定位每行第一个非NaN元素的位置，结合掩码判断全NaN的边界情况，逻辑简洁高效。

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