Определите массив с максимальным значением, поэлементно в numpy

Question

Я ищу вариант ответа для на этот вопрос

Учитывая несколько массивов (3 или более), мне нужно создать новый массив, где каждый элемент идентифицирует входные данные массив с максимальным значением.

например,

array0 = np.array([1, 2, 3])
array1 = np.array([0, 3, 4])
array2 = np.array([-1, 1, 1])

результирующий массив должен быть array([0,1,1])

np.maximum.reduce, похоже, не работает в этом случае , Я мог придумать только грубую силу для l oop

d = []
for i in range(len(array0)):
  k = 0
  if array1[i] > array0[i]:
    k = 1
  if array2[i] > array1[i] and array2[i] > array0[i]:
    k = 2
  d.append(k)

Есть ли более питонический / numpy способ сделать это?

Answer 1

Создайте новый массив из трех и возьмите argmax вдоль первой оси:

np.array([array0, array1, array2]).argmax(0)
# array([0, 1, 1])