numpy.searchsorted для нескольких экземпляров одной и той же записи - python

Question

У меня есть следующие переменные:

импортировать numpy как np

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

Я хочу вернуть записи gens, которые соответствуют каждому элементу p.

В идеале я хотел бы, чтобы он возвратил:

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]