CODEWITHSUNDEEP
pythonnumpy
Haden96
Haden96

Reputation: 47

Find element by row index and value

I need to find indexes of array by row index and value. 


     IN: arr1=[[0   5 ]
               [1   8 ]
               [2   3 ]]

         arr2=[[10  5   10 ]
               [8   10  10 ]
               [3   1   11 ]]
     OUT:
           [[0 1]
            [1 0]
            [2 0]]


     np.argwhere(arr2[arr1[:, 0], :] == arr1[:, 1])

It works in loop but I need faster way,I tried to do it in single argwhere but it returns an empty array.

      arr1 = np.array([[0,5],[1,8],[2,3]])

      arr2 = np.array([[10,5,10],[8,10,10],[3,1,11]])

      for (u,v) in arr1:
          result=np.argwhere(arr2[u, :] == v)

Upvotes: 1

Views: 123

Answers (2)

Divakar
Divakar

Reputation: 221684

You were simply needed to slice the second col of arr1 while keeping dimensions. One way to do so is with extending dims with np.newaxis/None after slicing (that reduces dims), hence -

np.argwhere(arr2[arr1[:, 0], :] == arr1[:, 1, None])

For a compact look, skip the trailing colon indexing -

np.argwhere(arr2[arr1[:,0]] == arr1[:,1,None])

More compact way, but one that forces copy of the second col of arr1 -

np.argwhere(arr2[arr1[:,0]] == arr1[:,[1]])

Note that np.argwhere gives us (row,col) pairs in a two-column array. With your loop-based solution you are basically extracting col indices. So, use the second column of the output from the posted code.

Upvotes: 2

Károly Szabó
Károly Szabó

Reputation: 1273

Some performance comparison:

%timeit np.argwhere(arr2[arr1[:, 0]] == arr1[:, 1, None])[:, 1]

8.82 µs ± 531 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

%timeit [np.argwhere(arr2[u, :] == v)[0,0] for (u,v) in arr1]

20.9 µs ± 2.74 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

Upvotes: 1

Related Questions