Индекс минимального значения выбранных элементов из массива питона

Question

С учетом массива numpy (a) и массива mask (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])]

Есть ли лучшие решения?Спасибо.

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() Истиной, а затем сделать все остальное.

Вот две функции для выполнения задачи:

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)