python - Find the row indexes of several values in a numpy array -

i have array x:

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

and wish find index of row of several values in array:

searched_values = np.array([[4, 2],                             [3, 3],                             [5, 6]]) 

for example result like:

[0,3,4] 

i have code doing this, think overly complicated:

x = np.array([[4,  2],               [9,  3],               [8,  5],               [3,  3],               [5,  6]])  searched_values = np.array([[4, 2],                             [3, 3],                             [5, 6]])  result = []  s in searched_values:     idx = np.argwhere([np.all((x-s)==0, axis=1)])[0][1]     result.append(idx)  print(result) 

i found this answer similar question works 1d arrays.

is there way want in simpler way?

approach #1

one approach use numpy broadcasting, -

np.where((x==searched_values[:,none]).all(-1))[1] 

approach #2

a memory efficient approach convert each row linear index equivalents , using np.in1d, -

dims = x.max(0)+1 out = np.where(np.in1d(np.ravel_multi_index(x.t,dims),\                     np.ravel_multi_index(searched_values.t,dims)))[0] 

approach #3

another memory efficient approach using np.searchsorted , same philosophy of converting linear index equivalents -

dims = x.max(0)+1 x1d = np.ravel_multi_index(x.t,dims) searched_valuesid = np.ravel_multi_index(searched_values.t,dims) sidx = x1d.argsort() out = sidx[np.searchsorted(x1d,searched_valuesid,sorter=sidx)] 

please note np.searchsorted method assumes there match each row searched_values in x.

---

how np.ravel_multi_index work?

this function gives linear index equivalent numbers. accepts 2d array of n-dimensional indices, set columns , shape of n-dimensional grid onto indices mapped , equivalent linear indices computed.

let's use inputs have problem @ hand. take case of input x , note first row of it. since, trying convert each row of x linear index equivalent , since np.ravel_multi_index assumes each column 1 indexing tuple, need transpose x before feeding function. since, number of elements per row in x in case 2, n-dimensional grid mapped onto 2d. 3 elements per row in x, had been 3d grid mapping , on.

to see how function compute linear indices, consider first row of x -

in [77]: x out[77]:  array([[4, 2],        [9, 3],        [8, 5],        [3, 3],        [5, 6]]) 

we have shape of n-dimensional grid dims -

in [78]: dims out[78]: array([10,  7]) 

let's create 2-dimensional grid see how mapping works , linear indices computed np.ravel_multi_index -

in [79]: out = np.zeros(dims,dtype=int)  in [80]: out out[80]:  array([[0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0]]) 

let's set first indexing tuple x, i.e. first row x grid -

in [81]: out[4,2] = 1  in [82]: out out[82]:  array([[0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 1, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0],        [0, 0, 0, 0, 0, 0, 0]]) 

now, see linear index equivalent of element set, let's flatten , use np.where detect 1.

in [83]: np.where(out.ravel())[0] out[83]: array([30]) 

this computed if row-major ordering taken account.

let's use np.ravel_multi_index , verify linear indices -

in [84]: np.ravel_multi_index(x.t,dims) out[84]: array([30, 66, 61, 24, 41]) 

thus, have linear indices corresponding each indexing tuple x, i.e. each row x.

choosing dimensions np.ravel_multi_index form unique linear indices

now, idea behind considering each row of x indexing tuple of n-dimensional grid , converting each such tuple scalar have unique scalars corresponding unique tuples, i.e. unique rows in x.

let's take @ x -

in [77]: x out[77]:  array([[4, 2],        [9, 3],        [8, 5],        [3, 3],        [5, 6]]) 

now, discussed in previous section, considering each row indexing tuple. within each such indexing tuple, first element represent first axis of n-dim grid, second element second axis of grid , on until last element of each row in x. in essence, each column represent 1 dimension or axis of grid. if map elements x onto same n-dim grid, need consider maximum stretch of each axis of such proposed n-dim grid. assuming dealing positive numbers in x, such stretch maximum of each column in x + 1. + 1 because python follows 0-based indexing. so, example x[1,0] == 9 map 10th row of proposed grid. similarly, x[4,1] == 6 go 7th column of grid.

so, our sample case, had -

in [7]: dims = x.max(axis=0) + 1 # or x.max(0) + 1  in [8]: dims out[8]: array([10,  7]) 

thus, need grid of @ least shape of (10,7) our sample case. more lengths along dimensions won't hurt , give unique linear indices too.

concluding remarks : 1 important thing noted here if have negative numbers in x, need add proper offsets along each column in x make indexing tuples positive numbers before using np.ravel_multi_index.