numpy.searchsorted 用于同一条目的多个实例 - python

Question

我有以下变量：

将 numpy 导入为 np

gens = np.array([2, 1, 2, 1, 0, 1, 2, 1, 2])
p = [0,1]

我想返回与p 的每个元素匹配的gens 的条目。

所以理想情况下我希望它返回：

result = [[4],[2,3,5,7],[0,2,6,8]] 
#[[where matched 0], [where matched 1], [the rest]]

--

到目前为止，我的尝试仅适用于一个变量：

indx = gens.argsort()
res = np.searchsorted(gens[indx], [0])
gens[res] #gives 4, which is the position of 0

但我尝试使用

indx = gens.argsort()
res = np.searchsorted(gens[indx], [1])
gens[res] #gives 1, which is the position of the first 1.

所以：

• 如何搜索出现多次的条目
• 如何搜索多个条目，每个条目都有多次出现？

Answer 1

您可以使用np.where

>>> np.where(gens == p[0])[0]
array([4])

>>> np.where(gens == p[1])[0]
array([1, 3, 5, 7])

>>> np.where((gens != p[0]) & (gens != p[1]))[0]
array([0, 2, 6, 8])

或者np.in1d和np.nonzero

>>> np.nonzero(np.in1d(gens, p[0]))[0]

>>> np.nonzero(np.in1d(gens, p[1]))[0]

>>> np.nonzero(~np.in1d(gens, p))[0]