查找值在给定集中的numpy向量的所有索引

Question

I am getting more and more used to numpy 's fancy indexing possibilities, but this time I hit an obstacle I cannot solve without resorting to ugly for loops.

My input is a pair of vectors, one large vector v and a smaller vector of indices e . What I want is to find all the indices i for which v[i] is equal to one of the values v[e[0]], v[e[1]],...v[e[n]] . At the moment, the code that does this for me (and it works) is

import numpy as np
v = np.array([0,0,0,0,1,1,1,2,2,2,2,2,2])
e=np.array([0,4])
#what I want to get is the vector [0,1,2,3,4,5,6].
values = v[e]
r = []
for i in range(n):
    if v[i] in values:
        r.append(i)

In the case when e is only one number, I am able to do this:

rr = np.arange(n)
r = v[rr] == v[e]

which is both nicer and quicker than a for loop. Is there a way of doing this when e is not a single number?

Answer 1

You could use where and in1d :

>>> v = np.array([0,0,0,0,1,1,1,2,2,2,2,2,2])
>>> e = [0,4]
>>> np.in1d(v, v[e])
array([ True,  True,  True,  True,  True,  True,  True, False, False,
       False, False, False, False], dtype=bool)
>>> np.where(np.in1d(v, v[e]))
(array([0, 1, 2, 3, 4, 5, 6]),)
>>> np.where(np.in1d(v, v[e]))[0]
array([0, 1, 2, 3, 4, 5, 6])