如何对含np.nan的numpy数组进行分段求和？

方法一：使用np.split结合np.nansum #

这种方法逻辑直观，适合处理小型数组：

import numpy as np

# 定义输入数组
arr = np.array([1, 2, 3, 4, np.nan, 2, 4, np.nan, 5, 1, 2, 1, np.nan, 10])

# 定位所有NaN的索引位置
nan_positions = np.where(np.isnan(arr))[0]
# 按NaN位置分割数组
segments = np.split(arr, nan_positions)

# 对每个分段求和（自动忽略分段末尾的NaN）
result = np.array([np.nansum(seg) for seg in segments])

print(result)
# 输出：array([10.,  6.,  9., 10.])

说明： #

• np.where(np.isnan(arr))精准定位所有NaN的位置，以此作为数组分割点。
• np.split将原数组拆分为多个子数组，每个子数组对应一段连续的非NaN元素（末尾可能附带NaN）。
• np.nansum会自动跳过子数组中的NaN，直接计算有效元素的总和。

方法二：分组求和（高效适配大型数组） #

这种方法利用分组标签批量计算，性能更优，适合处理大规模数据：

import numpy as np

arr = np.array([1, 2, 3, 4, np.nan, 2, 4, np.nan, 5, 1, 2, 1, np.nan, 10])

# 生成标记NaN位置的布尔掩码
mask = np.isnan(arr)
# 生成分组标签：每遇到一个NaN，后续元素的分组编号递增
groups = np.cumsum(mask)

# 筛选非NaN元素，按分组标签批量求和
result = np.bincount(groups[~mask], weights=arr[~mask])

print(result)
# 输出：array([10.,  6.,  9., 10.])

说明： #

• mask用布尔值标记原数组中的NaN位置。
• np.cumsum(mask)通过累加掩码值，为每一段连续非NaN元素分配唯一的分组编号。
• np.bincount根据分组编号，对对应位置的非NaN元素求和，直接得到各段的总和。

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