Searching a Numpy Array for the index of a subarray based on a subarray of the subarray

user_14933773 0
Viewed 106

I want to get the index of a 2d array which contains a specific array. In this case I want to know where in the array array the array [[4, 5], 6] is but only based on the inner most array [4, 5] so that I get its position even if instead of the six it would be an eight.

This is my code so far:

import numpy as np

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

print(np.where(array == [4, 5]))

but as an output I get:

(array([], dtype=int32), array([], dtype=int32))

and the output I want is the following:

(array([1], dtype=int32), array([0], dtype=int32))
numpy-ndarray multidimensional-array numpy arrays python
1 Answers

The issue is that you are working with a dtype object where your first numpy column contains list objects.

You can create a vectorized function to check each object individually.

f = np.vectorize(lambda x: x==[4,5])
idx = np.where(f(array))
idx

(array([1]), array([0]))

You can also use a list comprehension after flattening out the array and then checking each object against [4,5]. Then you could use np.where or just plain simple boolean check to get the index in that flat list, which you could unravel_index to get the position in the 2D array. (I am using np.where because you want to use it)

check = [i==[4,5] for i in array.ravel()]
np.unravel_index(*np.where(check), array.shape)

(array([1]), array([0]))
user_47559540
Related