0

I have a two equal sized 2d arrays for which I what to find the maximum value in the first and the corresponding value in the other.

I can get the index for the maximum in the first array by using numpy.argmax(). This gives me an array of indexes for which I want to find the value in the second 2d array corresponding to that index.

import numpy as np

a = np.array([[4,1,5,6],[6,8,1,2],[4,3,5,2],[5,6,2,8]])
b = np.array([[2,6,3,8],[3,5,8,3],[4,6,8,9],[6,2,4,7]])

index_max_a = np.argmax(a, axis=1)
print(index_max_a)

Which gives me:

array([3,1,2,3], dtype=int64)

Then my naive approach was to do:

b[:,index_max_a]

but this does not give the searched for list.

I can get the corresponding value to that index in b one by one as:

b[0,3]
b[1,1]
b[2,2]
b[3,3]

but can I do this in one swoop since I have the indexes in an array? I want to avoid using a for loop, since the arrays I am handling are quite large.

Ultimately a want to get a list which looks like (from this example):

array([8, 5, 8, 7])

It might be some syntax which I am missing. Thank you.

Improve this question
6
  • 1
    b[np.arange(index_max_a.size), index_max_a] Commented Oct 12, 2022 at 15:53
  • Just what I needed. Thank you very much. On a more technical note: What is the difference between inputting a list as the first argument compared to the ":" which I tried? Commented Oct 12, 2022 at 16:01
  • Does this answer your question? using integer as index for multidimensional numpy array (I am the author of the accepted answer.) Commented Oct 12, 2022 at 16:14
  • So in short: It is because the way the semicolon ":" indexes the rows is not the same way as calling each row manually, since there is some broadcasting performed in the semicolon operation? Commented Oct 12, 2022 at 16:23
  • 1
    It can also be solved with one line, but it requires you to manually expand the axis to meet the broadcast conditions: b[np.arange(b.shape[0])[:, None], np.arange(b.shape[1])[None], np.argmax(a, axis=2)] Commented Oct 13, 2022 at 2:29
Related questions
12224
How can I remove a specific item from an array in JavaScript?
4484
How can I find the index for a given item in a list?
5660
How can I access the index value in a 'for' loop?
Load 7 more related questions Show fewer related questions

0

Reset to default

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.