Find 1D array from 2D NumPy arrays?

Question

Suppose I have

[[array([x1, y1]), z1]
 [array([x2, y1]), z2]
 ......
 [array([xn, yn]), zn]
]

And I want to find the index of array([x5, y5]) . How can find effieciently using NumPy?

Answer 1

To start off, owing to the mixed data format, I don't think you can extract the arrays in a vectorized manner. Thus, you can use loop comprehension to extract the first element corresponding to the arrays from each list element as a 2D array. So, let's say A is the input list, we would have -

arr = np.vstack([a[0] for a in A])

Then, simply do the comparison in a vectorized fashion using NumPy's broadcasting feature , as it will broadcast that comparison along all the rows and look all matching rows with np.all(axis=1) . Finally, use np.flatnonzero to get the final indices. Thus, the final peace of the puzzle would be -

idx = np.flatnonzero((arr == search1D).all(1))

You can read up on the answers to this post to see other alternatives to get indices in such a 1D array searching in 2D array problem.

Sample run -

In [140]: A
Out[140]: 
[[array([3, 4]), 11],
 [array([2, 1]), 12],
 [array([4, 2]), 16],
 [array([2, 1]), 21]]

In [141]: search1D = [2,1]

In [142]: arr = np.vstack([a[0] for a in A]) # Extract 2D array

In [143]: arr
Out[143]: 
array([[3, 4],
       [2, 1],
       [4, 2],
       [2, 1]])

In [144]: np.flatnonzero((arr == search1D).all(1)) # Finally get indices
Out[144]: array([1, 3])