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

Question

给定一个 numpy 数组 (a) 和一个掩码数组 (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

更新:根据您的评论，要实现您想要的效果，您可以使用 np.nan_to_num在 a 上将无穷大替换为 a.dtype 可表示的最大有限浮点值，即 np.finfo(a.dtype).max 。请记住，此函数还会将 NaN 替换为 0，因此您可能需要将它们屏蔽掉或替换为其他值。

>>> 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

更新 2: 问题似乎是，如果屏蔽数组的所有未屏蔽值都是 inf，则 argmin 始终返回 0:

>>> 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

这是一个错误还是故意的？无论如何，为了处理这个问题，我们可以首先检查 np.isinf(a_masked).all() 是否为 True，然后执行其余操作。

这里有两个函数来实现这个任务:

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])]

恕我直言，OP提出的第二个版本看起来更好，更重要的是它更快:

N = 10000 
m = np.random.randint(2, size=N)
a = np.random.randint(N, size=N)*1.0
np.put(a, np.random.choice(range(N), N//2, replace=False), np.inf)

%timeit argmin_ma(a, m)
532 µs ± 70.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

%timeit argmin_ma2(a, m)
132 µs ± 6.61 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)