2

Is there any way to get the indices of several elements in a NumPy array at once?

For example:

import numpy as np
a = np.array([1, 2, 4])
b = np.array([1, 1, 3, 2, 4])

I would like to find the index of each element of a in b, namely: [0, 1, 3, 4].

Please Note:

  1. b has duplicated values, e.g. 1 here, methods for example in Getting the indices of several elements in a NumPy array at once would not work because it only find left-most or right-most index, not all indices. So using the method there would get [0, 3, 4] if left-most applied.
  2. I want to honour the order of the values in a, i.e. the first digits in the result is for the first value in a, and second few digits are for second value in a and so on, e.g. [0, 1] is for 1 in a, [3] is for 2 in a, and [4] is for 4 in a, so order in answer is [0, 1, 3, 4]
Shaido
  • 27,497
  • 23
  • 70
  • 73
Ridfield Cris
  • 75
  • 7
  • 1
    @Divakar: The solutions in the duplicate does not consider the condition that the `b` array can have duplicate values. The answers there does not really work in this more complex case. – Shaido May 15 '20 at 09:41
  • @Shaido-ReinstateMonica You need to read more on the accepted answer there. `Searchsorted` solution with `sorter` works. – Divakar May 15 '20 at 11:58
  • @Divakar: That solution will not consider any duplicates in `b` (note that `a` and `b` are swapped with the `A` and `B` in the linked question, i.e. `a=B` and `b=A` due to how the question is phrased). Using `searchsorted` with `sorter` for the example in the question will give `[0, 3, 4]` while OP wants `[0,1,3,4]`. – Shaido May 16 '20 at 04:36
  • 1
    @Shaido-ReinstateMonica Got it. reopened. – Divakar May 16 '20 at 06:07

1 Answers1

3

Thanks to my friend @rongkaizhang, I posted his beautiful answer here.

np.where(b == np.expand_dims(a, axis=1))[1]
Ridfield Cris
  • 75
  • 7