How to find the indices of an ndarray in another ndarray

Viewed 87

I have arrays of a = np.arange(21).reshape(7,3) and b = np.array([[3,4,5], [9,10,11]])

a = [[ 0  1  2]
     [ 3  4  5]
     [ 6  7  8]
     [ 9 10 11]
     [12 13 14]
     [15 16 17]
     [18 19 20]]

b = [[3  4  5]
     [9 10 11]]

I want to find the row numbers of b within a. So, I am looking to get 1 and 3 as my output. I know that for finding the indices, I can use the np.where() or np.argwhere(). However, I don't know exactly if we can use them for this problem or I have to use other functions. I tried c = np.argwhere(a == b) but it gives error.

1 Answers

You could use np.argwhere as follows:

import numpy as np

a = np.array([[0, 1, 2],
              [3, 4, 5],
              [6, 7, 8],
              [9, 10, 11],
              [12, 13, 14],
              [15, 16, 17],
              [18, 19, 20]])

b = np.array([[3, 4, 5],
              [9, 10, 11]])

res = np.argwhere(
    (a == b[:, None])  # compare all rows of a vs b
    .all(axis=2)  # find the ones where all the elements matches
)[:, 1]  # use argwhere to find indices, but only use the column indices

print(res)

Output

[1 3]

UPDATE For finding the missing ones do the following, I splitted the steps to make it easier to understand:

matches = (a == b[:, None]).all(axis=2)
print(matches)

res = np.argwhere(~matches.any(axis=1))[:, 0]
print(res)

Output

[[False  True False False False False False]
 [False False False False False False False]]
[1]

The first part of the output, shows two rows that correspond to the rows in b, as it can be seen the first row of b has a match in with the second row of a. The second row of, has no matches.

The second part of the output shows the result of applying argwhere for selecting the indices where there is not row of a matching one in b (~matches.any(axis=1)).

Related