CODEWITHSUNDEEP
pythonarraysnumpy
IssamLaradji
IssamLaradji

Reputation: 6865

np.where equivalent for multi-dimensional numpy arrays

Assume you have a numpy array as array([[5],[1,2],[5,6,7],[5],[5]]). Is there a function, such as np.where, that can be used to return all row indices where [5] is the row value? For example, in the array above, the returned values should be [0, 3, 4] indicating the [5] row numbers.

Please note that each row in the array can differ in length.

Thanks folks, you all deserve best answer, but i gave the green mark to the first one :)

Upvotes: 2

Views: 2637

Answers (3)

user2647939
user2647939

Reputation:

You could use the the list comprehension as here:

[i[0] for i,v in np.ndenumerate(a) if 5 in v]
#[0, 2, 3, 4]

Upvotes: 2

Saullo G. P. Castro
Saullo G. P. Castro

Reputation: 58885

If you check ndim of your array you will see that it is actually not a multi-dimensional array, but a 1d array of list objects.

You can use the following list comprehension to get the indices where 5 appears:

[i[0] for i,v in np.ndenumerate(a) if 5 in v]
#[0, 2, 3, 4]

Or the following list comprehension to get the indices where the list is exactly [5]:

[i[0] for i,v in np.ndenumerate(a) if v == [5]]
#[0, 3, 4]

Upvotes: 3

mattexx
mattexx

Reputation: 6606

This should do it:

[i[0] for i,v in np.ndenumerate(ar) if v == [5]]
=> [0, 3, 4]

Upvotes: 3

Related Questions