执行numpy“argany”或“argfirst”的方法？

Question

我想通过在其他维度中查找真实/匹配某些条件的 first 值来减少沿一个轴的多维数组，或者如果找不到匹配元素则使用默认值。到目前为止，我正在尝试通过一些简单的循环对 3D 数组执行此操作，如下所示：

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 有更好的内置方法来实现这一点吗？

Answer 1

我会将计划分为以下步骤：

1- 创建一个掩码来标识满足条件的位置。

2- 查找沿指定轴第一次满足条件的索引。

3- 使用这些索引从原始数组中收集相应的元素。

我会尝试这样的事情：

import numpy as np

def first(arr, condition, default):
    # Create a mask where the condition is met
    mask = condition(arr)
    
    # Initialize the output array with the default value
    out = np.full(arr.shape[1:], default)
    
    # Find the indices of the first occurrence of the condition being met along axis 0
    idx = np.argmax(mask, axis=0)
    
    # Create a mask to identify where no condition is met
    no_condition_met = ~mask.any(axis=0)
    
    # Use advanced indexing to gather the first occurrence where the condition is met
    out[~no_condition_met] = arr[idx[~no_condition_met], np.arange(arr.shape[1])[:, None], np.arange(arr.shape[2])]
    
    return out

# Example usage
arr = np.random.rand(4, 3, 3)
condition = lambda x: x > 0.5
default = -1

result = first(arr, condition, default)
print(result)