找到numpy向量的所有索引,其值在给定集合中

时间:2014-10-27 13:28:14

标签: python numpy

我越来越习惯于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不是一个数字时,有没有办法做到这一点?

1 个答案:

答案 0 :(得分:3)

您可以使用wherein1d

>>> 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])