如何在Numpy中批量校验多个索引区间内元素是否全为NaN

核心思路 #

利用前缀和数组快速计算任意区间的真值总和，通过比对总和与区间长度判断区间内是否全为True，完全规避逐区间循环，时间复杂度为线性，性能远高于原生Python循环。 #

实现代码 #

import numpy as np

# 边界处理：避免区间起止超出布尔数组的索引范围
max_valid_idx = len(is_element_from_another_array_nan)
indices_array = np.clip(indices_array, 0, max_valid_idx)

# 生成前缀和数组，首位补0以简化区间和计算逻辑
prefix_sum = np.cumsum(is_element_from_another_array_nan, dtype=np.int64)
prefix_sum = np.insert(prefix_sum, 0, 0)

# 批量计算所有区间的校验结果，统计符合条件的事件总数
interval_len = indices_array[:, 1] - indices_array[:, 0]
interval_truth_sum = prefix_sum[indices_array[:, 1]] - prefix_sum[indices_array[:, 0]]
counter = np.sum(interval_truth_sum == interval_len)

原理解释 #

布尔数组中True会被自动转换为1、False转换为0，因此区间内所有元素的和等于区间长度的充要条件就是区间内所有元素都是True。前缀和数组可以在O(1)时间复杂度内得到任意区间的元素和，整个方案仅需要两次线性遍历，即使区间数量达到十万、百万级也能快速处理，计算结果和你原有循环逻辑完全等价。 #

可选替代方案（仅适用于区间起始值严格递增的场景） #

如果你的indices_array第一列（区间起始索引）是严格递增的，也可以用np.logical_and.reduceat实现：

reduceat_res = np.logical_and.reduceat(is_element_from_another_array_nan, indices_array[:, 0])
# 过滤跨区间的冗余结果，仅取对应区间的校验值
counter = np.sum(reduceat_res[:len(indices_array)])

该方案性能和前缀和法接近，但兼容性较弱，更推荐使用前缀和方案。 #

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