Find index of multiple queries in a multidimensional numpy array

Question

I am looking for a way to find indices of an array of queries in a multidimensional array. For example:

arr = np.array([[17,  5, 19,  9],
   [18, 13,  3,  7],
   [ 8,  1,  4,  2],
   [ 8,  9,  7, 19],
   [ 6, 11,  8,  5],
   [11, 16, 13, 18],
   [ 0,  1,  2,  9],
   [ 1,  7,  4,  6]])

I can find indices for one query:

np.where(arr==1)
# (array([2, 6, 7]), array([1, 1, 0]))

Is there any numpy solution to do this for multiple values replacing the following for loop?

for q in queries:
    np.where(arr==q)

If both the array and queries were one-dimensional, I could use searchsorted as this answer but it doesn't work for multidimensional arrays.

shivsn · Accepted Answer · 2016-08-26 08:07:40Z

2

IIUC you may try this:

In[19]:np.where((arr==4)|(arr==5))
Out[19]: (array([0, 2, 4, 7], dtype=int64), array([1, 2, 3, 2], dtype=int64))

answered Aug 26, 2016 at 8:07

shivsn

7,9141 gold badge28 silver badges34 bronze badges

Sign up to request clarification or add additional context in comments.

Comments

Alexander · Accepted Answer · 2016-08-26 07:22:40Z

1

You can get the indices of each matching value by zipping the results of your where function and then using the * dereferencing operator.

arr = np.array([[17,  5, 19,  9],
                [18, 13,  3,  7],
                [ 8,  1,  4,  2],  # (2, 1)
                [ 8,  9,  7, 19],
                [ 6, 11,  8,  5],
                [11, 16, 13, 18],
                [ 0,  1,  2,  9],  # (6, 1)
                [ 1,  7,  4,  6]])  # (7, 0)

>>> zip(*np.where(arr == 1))
[(2, 1), (6, 1), (7, 0)]

I'm not sure what your intended output is, but you can use a dictionary comprehension to show the index location for a given set of numbers, e.g.:

>>> {n: zip(*np.where(arr == n)) for n in range(5)}
{0: [(6, 0)],
 1: [(2, 1), (6, 1), (7, 0)],
 2: [(2, 3), (6, 2)],
 3: [(1, 2)],
 4: [(2, 2), (7, 2)]}

edited Aug 26, 2016 at 7:22

answered Aug 26, 2016 at 6:25

Alexander

111k32 gold badges212 silver badges208 bronze badges

3 Comments

Mahdi Over a year ago

Thanks but I wanted to get np.where(arr == q) for different values of q!

Alexander Over a year ago

Do you want the indices for each value of q or for all of the values of q?

user1121588 Over a year ago

>>> ix = np.in1d(arr.ravel(), [1,8,18]).reshape(arr.shape) ... >>> found = np.array(list(zip(np.where(ix)))) ... list is need in python 3.

Collectives™ on Stack Overflow

Find index of multiple queries in a multidimensional numpy array

2 Answers 2

Comments

3 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

Comments

3 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related