Ad

Extracting All Second Column Values From A Ndarray For A Given First Column Value

- 1 answer

I have a numpy ndarray where first column is user id and second column is some product id. What would be the fastest way to get all product ids for a given user id?

I've been going through the numpy doc and this handbook (https://jakevdp.github.io/PythonDataScienceHandbook/02.02-the-basics-of-numpy-arrays.html) as well but I had no luck.

Say we have this array:

test = [[0, 1], [0, 20], [0, 30], [1, 11], [1, 23], [1, 45]]

My goal is to get a function like this:

get_product_ids(0)

>> [1, 20, 30]
Ad

Answer

This can be achieved in such a simple way

test = np.array(test)

def get_product_id(ind):
    mask = test[:, 0] == ind
    return test[:, 1][mask]
Ad
source: stackoverflow.com
Ad