如何高效从沿第三维度堆叠的numpy数组中提取首个非NaN值构建2D数组？

需求说明 #

现有n个沿第三维度堆叠的2D NumPy数组，每个数组包含np.nan值。需要生成一个新的2D数组，其中每个元素为沿第三维度的首个非np.nan值（注：示例中c数组优先级最高，即从最后一个堆叠的数组开始往前查找第一个非np.nan值）。

原循环实现 #

此前通过循环实现了需求，代码如下：

import numpy as np
a = np.array([[1,2,3],
              [4,5,np.nan],
              [7,8,9],
              [10,11,12]])
b = np.array([[np.nan,21,np.nan],
              [23,24,25],
              [26,27,28],
              [29, 30,31]])
c = np.array([[np.nan,np.nan,np.nan],
              [43,np.nan,np.nan],
              [46,np.nan,48],
              [49, 50, 51]])
stack = np.stack([a, b, c], axis=0)

for d in [2,1]:
    if d == 2:
        fill = stack[d, :,:]
    arr = np.where(np.isnan(stack[d,:,:]), stack[d-1,:,:], fill)
    fill = arr

# 预期结果
>>> arr
array([[ 1., 21.,  3.],
       [43., 24., 25.],
       [46., 27., 48.],
       [49., 50., 51.]])

无循环高效解法 #

方法1：链式np.where（适合少量数组） #

利用np.where的嵌套逻辑，从高优先级到低优先级依次覆盖NaN值，代码简洁直观：

import numpy as np

# 定义输入数组（同原示例）
a = np.array([[1,2,3],
              [4,5,np.nan],
              [7,8,9],
              [10,11,12]])
b = np.array([[np.nan,21,np.nan],
              [23,24,25],
              [26,27,28],
              [29, 30,31]])
c = np.array([[np.nan,np.nan,np.nan],
              [43,np.nan,np.nan],
              [46,np.nan,48],
              [49, 50, 51]])

# 直接链式判断：优先取c，c为NaN则取b，b也为NaN则取a
result = np.where(np.isnan(c), np.where(np.isnan(b), a, b), c)

# 输出结果
>>> result
array([[ 1., 21.,  3.],
       [43., 24., 25.],
       [46., 27., 48.],
       [49., 50., 51.]])

方法2：通用索引提取（适合任意数量数组） #

当堆叠的数组数量较多时，手动嵌套np.where会变得繁琐，可通过索引定位实现通用解法：

import numpy as np

# 定义输入数组并堆叠
a = np.array([[1,2,3],
              [4,5,np.nan],
              [7,8,9],
              [10,11,12]])
b = np.array([[np.nan,21,np.nan],
              [23,24,25],
              [26,27,28],
              [29, 30,31]])
c = np.array([[np.nan,np.nan,np.nan],
              [43,np.nan,np.nan],
              [46,np.nan,48],
              [49, 50, 51]])
stack = np.stack([a, b, c], axis=0)

# 1. 生成非NaN值的掩码
non_nan_mask = ~np.isnan(stack)
# 2. 反转堆叠维度，找到每个位置第一个非NaN值的索引（反转后argmax取第一个True的位置）
first_valid_rev_idx = np.argmax(non_nan_mask[::-1], axis=0)
# 3. 转换为原堆叠数组的索引
first_valid_idx = stack.shape[0] - 1 - first_valid_rev_idx
# 4. 提取对应位置的值
result = np.take_along_axis(stack, first_valid_idx[np.newaxis, ...], axis=0).squeeze(axis=0)

# 输出结果
>>> result
array([[ 1., 21.,  3.],
       [43., 24., 25.],
       [46., 27., 48.],
       [49., 50., 51.]])

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