在numpy中识别具有最大值的数组，元素明智

Question

我正在寻找this 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 在这种情况下似乎不起作用。 我只能想出一个暴力循环

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)

有没有更 Pythonic/numpy 的方式来做到这一点？

Answer 1

从三个中构建一个新数组，并沿第一个轴取argmax：

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