CODEWITHSUNDEEP
pythonnumpyvectorization
Diego Aguado
Diego Aguado

Reputation: 1596

How to use numpy where for several possible values?

Consider a numpy ndarray called picks_user with shape picks_user.shape = (2016,3). The 'columns' represent the variables user, item and count in that order. The 'rows' represent observations.

When performing:

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

the result is another numpy ndarray with the users that have select item 2711.

Say that target_users has shape target_users.shape = (14,). I want to use this array to get all items picked by the target users, something like the following (which doesnt work):

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

This could be equivalent to:

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

How can I achieve this in a vectorized way, no for loop?

UPDATE

Consider the following example:

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]

In this case target_users = [1 5]

The vector which I want to get is:

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

Upvotes: 1

Views: 84

Answers (1)

Irshad Bhat
Irshad Bhat

Reputation: 8709

You can use 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]])

Upvotes: 3

Related Questions