Find location of pair of elements in two arrays in numpy

Question

I have two numpy arrays x and y

Suppose x = [0, 1, 1, 1, 3, 4, 5, 5, 5] and y = [0, 2, 3, 4, 2, 1, 3, 4, 5]

The length of both arrays is the same and the coordinate pair I am looking for definitely exists in the array.

How can I find the index of (a, b) in these arrays, where a is an element in x and b is the corresponding element in y. For example, the index of (1, 4) would be 3: the elements at index 3 of x and y are 1 and 4 respectively.

Answer 1

You could use numpy.where combined with numpy.logical_and if you want a purely numpy solution:

In [16]: import numpy as np

In [17]: x = np.array([0, 1, 1, 1, 3, 4, 5, 5, 5])

In [18]: y = np.array([0, 2, 3, 4, 2, 1, 3, 4, 5])

In [19]: np.where(np.logical_and(x == 1, y == 4))[0]
Out[19]: array([3], dtype=int64)

numpy.logical_and allows you to element-wise perform a logical AND operation between two numpy arrays. What we're doing here is determining which locations contain both the x values being 1 and the y values being 4 in the same corresponding locations. Those locations that satisfy this are True. numpy.where determines the locations in the array where this condition is satisfied. numpy.where actually returns both row and column locations of where the condition is True separately as a tuple of two elements, but as we are only considered with one dimension, only the first tuple is valid which is why we immediately index the first element of the tuple.

The output is a numpy array of locations where the condition is valid. You can even go further and coerce the output to be a list of indices to make things neater and/or if it is required (thanks @EddoHintoso):

In [20]: list(np.where(np.logical_and(x == 1, y == 4))[0])
Out[20]: [3]

Answer 2

You could compare your first array with the first value, second array with the second value and then find where both True. Then you could get that True with argmax which will give you first index of the first True occurence:

x = np.array([0, 1, 1, 1, 3, 4, 5, 5, 5])
y = np.array([0, 2, 3, 4, 2, 1, 3, 4, 5])
idx = ((x == 1) & (y == 4)).argmax()

In [35]: idx
Out[35]: 3

In [36]: x == 1
Out[36]: array([False,  True,  True,  True, False, False, False, False, False], dtype=bool)

In [37]: y == 4
Out[37]: array([False, False, False,  True, False, False, False,  True, False], dtype=bool)

If you could have multiple occurence you could use following using nonzero:

idx_list = ((x == 1) & (y == 4))
idx = idx_list.nonzero()[0]

In [51]: idx
Out[51]: array([3], dtype=int64)

Or if you need list of indices:

In [57]: idx_list.nonzero()[0].tolist()
Out[57]: [3]

You could do that in one line with:

idx = ((x == 1) & (y == 4)).nonzero()[0]