时间戳

首页 > 解决方案 > python numpy数组中选定元素的最小值的索引

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

问题描述

给定一个numpy数组（a）和一个掩码数组（m），我怎样才能得到最小值的索引？例如，如果a = [3, 2, 4, 5]和m = [1, 0, 0, 1]，则预期答案为 1，因为在第 2 和第 3 个元素中，最小值是第 2 个元素（索引 1）。我的解决方案（似乎很尴尬）：

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

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

标签： pythonnumpy

解决方案

使用掩码数组怎么样？

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 ona将无穷大替换为可表示的最大有限浮点值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)

推荐阅读