CODEWITHSUNDEEP
pythonarraysnumpy
Carlos Eduardo Corpus
Carlos Eduardo Corpus

Reputation: 349

Getting the index of a 2D numpy array according to a 1D array

I'm working with numpy arrays I hit a roadblock, maybe you can help me. So I have a 2D numpy array, each array in this 2D array has a max value, what I need is get the index of the max value but not of the 2D array, instead, use the index like it was a 1D array, I know I can use numpy.argmax to find the index, but the index it's according to the 2D array not an 1D array, maybe an example would be better:

import numpy as np

arr =  np.array([[512, 523, 491], 
                 [512, 531, 495]])

index = np.argmax(arr, axis = 1)
index2 = np.argmax(arr)

print(index)
print(index2) 

index = [1, 1] 
index2 = 4

I mean it does what it's supposed to do, but what if I need the index like it was a 1D array? This is my desired output:

index = [1, 4]

So the first max value is 523 so the first index is 1, and the second max value is 531 so the second index is 4 according like it was a 1D array. Maybe this is a newbie question, but I'm not sure how to do this and this is just an example, the arr array can be bigger, so, any help will be appreciated, thank you!

Upvotes: 1

Views: 800

Answers (3)

Ibrahim Qasim
Ibrahim Qasim

Reputation: 21

You could add offsets to each element, given that you know the dimensions of the 2D array.

e.g.

index = [1, 1]
offset = [0, 3]

index + offset
>>> [1, 4]

You can get the offsets quite easily by doing np.arange(len(index))*size. Where size is the width of the 2D array.

Upvotes: 1

RichieV
RichieV

Reputation: 5183

The first line is actually what you wanted. It returns the index of the max in each nested array, i.e. 531 is in index 1 of the second nested array.

You can use that with advanced indexing as

>>> print(arr[np.arange(arr.shape[0]), index])
[523, 531]

Or change index to a flat index

flat_index = index + np.arange(0, arr.size, arr.shape[1])
# flat_index == [1, 4]

Upvotes: 2

David Kaftan
David Kaftan

Reputation: 2174

The easiest thing I could think to do is add an offset to index.

offset = np.arange(0,arr.shape[0]*arr.shape[1], arr.shape[1])
index1D = offset + index

for context, np.arange takes 3 args: start, stop, and stride. so the offset is just adding the total number of elements before a particular row to the index.

Upvotes: 3

Related Questions