np.where equivalent for multi-dimensional numpy arrays

Question

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 :)

Answer 1

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]

Answer 2

This should do it:

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

Answer 3

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]