Question

给出一个numpy数组（a）和一个mask数组（m），如何获得最小值的索引？例如，如果a = [3, 2, 4, 5]和m = [1, 0, 0, 1]的期望答案为1，因为在第二个和第三个元素中，最小值是第二个元素（索引1）。我的解决方案（似乎很尴尬）：

index = np.where(m == 0)[0]
point = index[np.argmin(a[index])]

有更好的解决方案吗？谢谢。

Answer 1

使用遮罩数组怎么办？

np.ma.array(a, mask=m).argmin()

示例：

>>> a = [0, np.inf, 1, 2]
>>> m = [1, 0, 0, 1]
>>> np.ma.array(a, mask=m).argmin()
2

更新：根据您的评论，要实现所需的目标，您可以在a上使用np.nan_to_num来将无限数替换为{{ 1}}，即a.dtype。请记住，此功能还将NaN替换为0，因此您可能要屏蔽掉它们或将其替换为其他值。

np.finfo(a.dtype).max

更新2：问题似乎在于，如果掩码数组的所有未掩码值都是inf的，那么>>> a = [np.inf, np.inf, np.inf] >>> m = [1, 0, 0] >>> a_masked = np.ma.array(np.nan_to_num(a), mask=m) >>> a_masked.argmin() 1总是返回0：

argmin

这是错误还是故意的？无论如何，要解决此问题，我们可以先检查>>> m = [1, 1, 0, 1, 0] >>> a = [10, 9, np.inf, 8, np.inf] >>> a_masked = np.ma.array(a, mask=m) >>> a_masked.argmin() 0是否为True，然后再进行其余操作。

以下是完成任务的两个功能：

np.isinf(a_masked).all()

OP提出的第二个版本看起来更好，更重要的是它更快：

def argmin_ma(a, m):
    if np.all(m):
        return None    
    a_masked = np.ma.array(a, mask=m)    
    if np.isinf(a_masked).all():    
        #a_masked = np.ma.array(np.nan_to_num(a), mask=m)
        #return a_masked.argmin()
        return np.argmin(m)
    return a_masked.argmin()

def argmin_ma2(a, m):
    if np.all(m):
        return None 
    a = np.asarray(a)
    m = np.aasrray(m)
    index = np.where(m == 0)[0]
    return index[np.argmin(a[index])]

python numpy数组中选定元素的最小值的索引

1 个答案: