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

Question

我越来越习惯于numpy的幻想索引可能性，但是这次我遇到了一个障碍，如果不采用丑陋的循环，我就无法解决这个问题。

我的输入是一对向量，一个大向量v和一个较小的向量索引e。我想要的是找到i的所有指数，其中v[i]等于值v[e[0]], v[e[1]],...v[e[n]]之一。目前，这是否对我（和它的作品）的代码是

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) 

在这种情况下，当e是只有一个号码，我能做到这一点：

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

这是比for循环更好更快。当e不是单个数字时，是否有这种方法？

Answer 1

你可以使用where和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])