CODEWITHSUNDEEP
pythonnumpy
SamGhatak
SamGhatak

Reputation: 1493

Find the row-wise index of a list of values from numpy matrix

Is there a smart way to find the index of a list of values in a matrix using numpy? We can always do it in iterative way, but is there any quick and fast way available?

We have a matrix as:

[[1,2,3],[4,5,6],[7,8,9],[10,11,12]]

and a list:

[3,4,7,11]

What we need here is the index of 3,4,7 and 11 in corresponding rows of the matrix, i.e. [2,0,0,1].

It is obviously a simple code to write in iterative approach, but we are looking for any ready implementation.

Thanks in advance.

Upvotes: 3

Views: 2017

Answers (3)

Sohaib Farooqi
Sohaib Farooqi

Reputation: 5666

You can use numpy.argwhere to get you desired indices

import numpy as np

arr = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])
search_list = [3,4,7,11]
index_arr = [np.argwhere(arr==x).flatten()[1] for x in search_list]
#index_arr [2, 0, 0, 1]

Upvotes: 1

Paul Panzer
Paul Panzer

Reputation: 53089

Here is a duplicate-safe version of @Akavall's solution. If there are multiple occurrences of a value the first index is returned:

a = [[1,2,3],[4,5,6],[7,8,9],[10,11,12]]
b = [3,4,7,11]

# introduce a repeated value 
a[1][1] = a[1][0]
a
# [[1, 2, 3], [4, 4, 6], [7, 8, 9], [10, 11, 12]]

i, j = np.where(a == np.c_[b])
j[i.searchsorted(range(len(b)))]
# array([2, 0, 0, 1])

Upvotes: 1

Akavall
Akavall

Reputation: 86306

For your example this works:

In [17]: import numpy as np

In [18]: a = np.array([[1,2,3],[4,5,6],[7,8,9],[10,11,12]])

In [19]: l = [3,4,7,11]

In [20]: np.where(a == np.array(l)[:, None])[1]
Out[20]: array([2, 0, 0, 1])

Upvotes: 2

Related Questions