CODEWITHSUNDEEP
pythonarraysnumpy
k3it
k3it

Reputation: 2730

Identify an array with the maximum value, element wise in numpy

I'm looking for a variation on the answer for this question

Given multiple arrays (3 or more), i need to construct a new array where each element identifies the input array which had the maximum value.

eg

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

the resulting array should be array([0,1,1])

np.maximum.reduce does not seem to work in the case. I could only come up with a brute force for loop

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)

Is there a more pythonic/numpy way to do this?

Upvotes: 0

Views: 92

Answers (1)

yatu
yatu

Reputation: 88226

Build a new array from the three and take the argmax along the first axis:

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

Upvotes: 1

Related Questions