Создание нового массива путем сравнения массивов

Question

Можно ли создать новый массив, сравнив массивы с numpy?

У меня есть 3 массива (A1, A2, A3). Как найти все индексы, где A1 == 2 and A2 > A3 и записать там значение 5 в новом массиве?

У меня есть этот MatLab код, который делает это:

index = find(A1==2 & A2>A3); 
new_array(index) = 5; 

я нашел putmask и logical_and, но не уверен, если это правильные инструменты и как использовать его в моем случае. Благодаря!

Answer 1

В приведенном ниже коде используется &, чтобы связать условные обозначения. В таких случаях, когда A1 == 2 и A2 > A3 оба True, то index массив будет True

import numpy as np 

A1 = np.array([0, 0, 2, 2, 4, 4]) 
A2 = np.arange(len(A1)) 
A3 = np.ones(len(A1))*3 

new_array = np.zeros(len(A1)) 

index = (A1 == 2) & (A2 > A3) 

new_array[index] = 5 
# array([ 0., 0., 5., 5., 0., 0.]) 

Вы можете использовать np.logical_and конечно. Но это ограничивает вас только двумя условными выражениями, в то время как вы можете эффективно целить столько, сколько хотите, используя &.

Answer 2

Вы можете использовать функцию np.where

import numpy as np 

A1 = np.array([0, 0, 2, 2, 4, 4]) 
A2 = np.arange(len(A1)) 
A3 = np.ones(len(A1))*3 

out = np.where((A1 == 2) & (A2 >= A3), 5, 0) 

или более простой

out = ((A1 == 2) & (A2 >= A3))*5