Извлечение всех значений второго столбца из ndarray для заданного значения первого столбца

Question

У меня есть nundy ndarray, где первый столбец - идентификатор пользователя, а второй столбец - идентификатор продукта. Каков был бы самый быстрый способ получить все идентификаторы продуктов для данного идентификатора пользователя?

Я просматривал документацию и этот справочник (также https://jakevdp.github.io/PythonDataScienceHandbook/02.02-the-basics-of-numpy-arrays.html), но мне не повезло.

Скажем, у нас есть этот массив:

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

Моя цель - получить такую ​​функцию:

get_product_ids(0)

>> [1, 20, 30]

Answer 1

Как насчет np.where и некоторой нарезки?

def get_prod_id(arr, user_id):
    return arr[np.where(arr[:,0] == user_id)[0],1]

test = np.array([[0, 1],[0, 20],[0, 30],[1, 11],[1, 23],[1, 45]])

get_prod_id(test, 1)
Out[32]: array([11, 23, 45])

get_prod_id(test, 0)
Out[33]: array([ 1, 20, 30])

... хотя это делает хороший 1-строчный - в отношении скорости выполнения, если вы бы не работали сnumpy arrays, чистое решение Python, подобное предложенному Amit Yadav , скорее всего, быстрее.

Answer 2

Это может быть достигнуто таким простым способом

test = np.array(test)

def get_product_id(ind):
    mask = test[:, 0] == ind
    return test[:, 1][mask]

Answer 3

def get_product_ids(user_id, user_products):
    products = []

    # In one line using list comprehension
    products.extend([sub_list[1] for sub_list in user_products if sub_list[0] == user_id])

    # if above code is too complicated, used the below three lines instead
    # for sub_list in user_products:
    #    if sub_list[0] == user_id:
    #        products.append(sub_list[1])

    return products

user_products = [[0,1],[0,20],[0,30],[1,11],[1,23],[1,45]]
get_product_ids(0, user_products)

Дайте выше попробовать. Я предоставил два способа сделать то же самое. Используйте то, что вы предпочитаете.

Сложность где-то рядом с O (n), которая должна быть достаточно хорошей и быстрой.

Если вы собираетесь использовать эту функцию много раз, я бы посоветовалнаписание другой функции, которая возвращала бы словарь с user_id в качестве ключей и product_ids в качестве значений. Это сделало бы всю работу по-настоящему эффективной, чем перебор полного списка каждый раз.

С вашим примером ввода этот словарь будет выглядеть примерно так:

{ 0: [1,20,30], 1: [11,23,45] }