Find Indices Of Numpy Array Based On Values In Another Numpy Array

I want to find the indices in a larger array if they match the values of a different, smaller array. Something like new_array below: import numpy as np summed_rows = np.random.rand

Solution 1:

Taking into account the proposed options on the comments, and adding an extra option with numpy's in1d option:

>>>import numpy as np>>>summed_rows = np.random.randint(low=1, high=14, size=9999)>>>common_sums = np.array([7,10,13])>>>ind_1 = (summed_rows==common_sums[:,None]).any(0).nonzero()[0]   # Option of @Brenlla>>>ind_2 = np.where(summed_rows == common_sums[:, None])[1]   # Option of @Ravi Sharma>>>ind_3 = np.arange(summed_rows.shape[0])[np.in1d(summed_rows, common_sums)]>>>ind_4 = np.where(np.in1d(summed_rows, common_sums))[0]>>>ind_5 = np.where(np.isin(summed_rows, common_sums))[0]   # Option of @jdehesa>>>np.array_equal(np.sort(ind_1), np.sort(ind_2))
True
>>>np.array_equal(np.sort(ind_1), np.sort(ind_3))
True
>>>np.array_equal(np.sort(ind_1), np.sort(ind_4))
True
>>>np.array_equal(np.sort(ind_1), np.sort(ind_5))
True

If you time it, you can see that all of them are quite similar, but @Brenlla's option is the fastest one

python -m timeit -s 'import numpy as np; np.random.seed(0); a = np.random.randint(low=1, high=14, size=9999); b = np.array([7,10,13])' 'ind_1 = (a==b[:,None]).any(0).nonzero()[0]'
10000 loops, best of 3: 52.7 usec per loop

python -m timeit -s 'import numpy as np; np.random.seed(0); a = np.random.randint(low=1, high=14, size=9999); b = np.array([7,10,13])' 'ind_2 = np.where(a == b[:, None])[1]'
10000 loops, best of 3: 191 usec per loop

python -m timeit -s 'import numpy as np; np.random.seed(0); a = np.random.randint(low=1, high=14, size=9999); b = np.array([7,10,13])' 'ind_3 = np.arange(a.shape[0])[np.in1d(a, b)]'
10000 loops, best of 3: 103 usec per loop

python -m timeit -s 'import numpy as np; np.random.seed(0); a = np.random.randint(low=1, high=14, size=9999); b = np.array([7,10,13])' 'ind_4 = np.where(np.in1d(a, b))[0]'
10000 loops, best of 3: 63 usec per loo

python -m timeit -s 'import numpy as np; np.random.seed(0); a = np.random.randint(low=1, high=14, size=9999); b = np.array([7,10,13])' 'ind_5 = np.where(np.isin(a, b))[0]'
10000 loops, best of 3: 67.1 usec per loop

Solution 2:

Use np.isin:

import numpy as np
summed_rows = np.random.randint(low=1, high=14, size=9999)
common_sums = np.array([7, 10, 13])
new_array = np.where(np.isin(summed_rows, common_sums))