如何用Numpy高效实现多维数组的argfirst/argany功能？

我需要沿指定轴约减多维数组，提取其他维度中首个满足指定条件的值；如果该维度下没有匹配元素，则使用默认值填充。目前我用循环实现了3D数组的该逻辑，但代码不够优雅、效率低，还没法轻松适配任意维度的数组。循环实现的代码如下：

import numpy as np

def first(arr, condition, default):
    out = default.copy()

    for u in range(arr.shape[1]):
        for e in range(arr.shape[2]):
            (idx,) = np.nonzero(condition[:, u, e])
            if len(idx):
                out[u, e] = arr[idx[0], u, e]

    return out

请问Numpy有没有更简洁的内置方法或向量化实现方式？

补充示例 #

测试数据：

test_arr = np.array(
    [
        [np.nan, 1, 2, 3, 4, np.nan, 6, 7, 8, 9],
        [np.nan, np.nan, 3, 4, 5, 6, 7, 8, 9, np.nan],
        [2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
        np.full((10,), np.nan),
    ]
)

test_default = np.array([5, 5, 5, 5])
test_cond = ~np.isnan(test_arr)

调用first(test_arr, test_cond, test_default)的预期输出：

[1.0, 3.0, 2.0, 5.0]

可以用Numpy的向量化操作实现通用版本，无需循环，且支持任意维度数组：

通用实现代码 #

import numpy as np

def first_vectorized(arr, condition, axis=0, default=None):
    # 获取沿目标轴的首个满足条件的索引
    first_idx = np.argmax(condition, axis=axis)
    # 构造用于提取值的索引数组（保持维度匹配）
    idx_expanded = np.expand_dims(first_idx, axis=axis)
    # 提取对应位置的值
    result = np.take_along_axis(arr, idx_expanded, axis=axis).squeeze(axis=axis)
    # 找出无匹配元素的位置
    no_match_mask = ~np.any(condition, axis=axis)
    # 替换为默认值
    if default is not None:
        if isinstance(default, np.ndarray):
            result[no_match_mask] = default[no_match_mask]
        else:
            result[no_match_mask] = default
    return result

验证示例 #

用你的测试数据验证（注意这里指定axis=1，因为测试数组是4行10列，要沿列轴找首个非NaN值）：

result = first_vectorized(test_arr, test_cond, axis=1, default=test_default)
print(result)
# 输出：array([1., 3., 2., 5.])

关键逻辑说明 #

• np.argmax(condition, axis=axis)：布尔数组中True等价于1，False等价于0，argmax会返回沿指定轴第一个最大值（即首个True）的索引；若该轴全为False，则返回0。
• np.take_along_axis：根据构造的索引数组直接提取对应值，完全替代循环遍历。
• np.any(condition, axis=axis)：判断每个非目标轴维度是否存在满足条件的元素，生成掩码后替换为默认值。
• 该实现支持任意维度数组，只需调整axis参数即可，无需修改核心逻辑。

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