Как использовать numpy, где для нескольких возможных значений?

Question

Рассматривают numpy ndarray под названием picks_user с формой picks_user.shape = (2016,3). «Столбцы» представляют переменные user, item и count в этом порядке. «Строки» представляют собой наблюдения.

При выполнении:

target_users = picks_user[np.where(picks_user[:,1]== 2711)][:,0]

результатом является еще одним NumPy ndarray с пользователями, которые имеют выберите пункт 2711.

Говорят, что target_users имеет форму target_users.shape = (14,). Я хочу использовать этот массив, чтобы получить все элементы подобран целевых пользователей, что-то вроде следующего (который не работает):

picks_user[np.where(picks_user[:,1] == target_users)] 

Это может быть эквивалентно:

for element in target_users: 
    picks_user[np.where(picks_user[:,1] == element] 

Как можно достичь это векторизованным способом, нет для цикла?

UPDATE

Рассмотрим следующий пример:

a = np.array([ [1,10,1],[2,11,1],[3,12,1],[4,13,1],[5,10,1],[2,13,1],[1,11,1],[5,16,1]]) 
target_users = a[np.where(a[:,1]==10)][:,0] 

В этом случае target_users = [1 5]

вектор, который я хочу, чтобы это:

[[1,10,1],[5,10,1],[1,11,1],[5,16,1]] 

Answer 1

Вы можете использовать np.in1d as:

>>> picks_user = np.random.randint(0,10, (10,3)) 
>>> picks_user 
array([[7, 8, 7], 
     [6, 0, 9], 
     [5, 6, 7], 
     [6, 7, 3], 
     [0, 1, 3], 
     [8, 7, 5], 
     [2, 6, 6], 
     [7, 9, 8], 
     [1, 7, 1], 
     [9, 8, 4]]) 
>>> target_users = np.array([1,7,8]) 
>>> picks_user[np.in1d(picks_user[:,1], target_users)] 
array([[7, 8, 7], 
     [6, 7, 3], 
     [0, 1, 3], 
     [8, 7, 5], 
     [1, 7, 1], 
     [9, 8, 4]])