Trying to retrieve index from an list using another list in python without pandas-CodePudding

So I have two 1d array, one being a subset of the other. Say, one array is called main with length 600, and the other array called subset with length 230.

I'm trying to find the specific index of the main array that corresponds to the subset array.

Say, I have

array([0.3,0.2,0.7,0.9]) as my main array and the subset after splicing giving ([0.3,0.9]). I want to find which index from the main does 0.3 and 0.9 come from (0 and 3) based on the subset. So I'm pretty much only interested in the index from the main array.

This is my recent attempt on solving this:

index = []
for i in subset:
    inf = np.where(main == subset[i])
    index.append(inf)

I get an error on the inf line due to indexerror:

IndexError: only integers, slices (:), ellipsis (...), numpy.newaxis (None) and integer or boolean arrays are valid indices

I know if I had pandas, I would have no issues finding it. But i'm restricted into using it, so I have to find a way doing this with numpy, which is faster anyways. Any assistance is appreciated!

CodePudding user response：

import numpy as np

main = np.array([0.3, 0.2,0.7,0.9])
subset = np.array([0.3, 0.9])

index = []
for i in subset:
    indices = np.where(main == i)[0]
    index.extend(indices)

print(index)