Tengo un ndarray numpy donde la primera columna es la identificación del usuario y la segunda columna es la identificación del producto. ¿Cuál sería la forma más rápida de obtener todos los identificadores de productos para una identificación de usuario determinada?

He estado revisando el documento numpy y este manual (https://jakevdp.github.io/PythonDataScienceHandbook/02.02-the-basics-of-numpy-arrays.html) también, pero no tuve suerte.

Digamos que tenemos esta matriz:

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

Mi objetivo es obtener una función como esta:

get_product_ids(0)

>> [1, 20, 30]
0
user11203229 6 oct. 2019 a las 12:10

3 respuestas

La mejor respuesta

Esto se puede lograr de una manera tan simple

test = np.array(test)

def get_product_id(ind):
    mask = test[:, 0] == ind
    return test[:, 1][mask]
0
meTchaikovsky 6 oct. 2019 a las 12:27

¿Qué pasa con np.where y algunas rebanadas?

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])

... aunque esto hace un buen 1-liner - con respecto a la velocidad de ejecución, si no está trabajando con matrices numpy, una solución pura de Python como la propuesta por Amit Yadav probablemente sea más rápido.

0
MrFuppes 6 oct. 2019 a las 09:58
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)

Prueba más arriba. He proporcionado dos formas de hacer lo mismo. Use lo que prefiera.

La complejidad está en algún lugar cerca de O (n), que debería estar bien y lo suficientemente rápido.

Si va a utilizar esta función muchas veces, le sugiero que escriba otra función que devuelva un diccionario con user_id como claves y product_ids como valores. Eso haría que la operación general sea realmente eficiente que iterar sobre la lista completa cada vez.

Con su entrada de muestra, ese diccionario se vería como a continuación:

{ 0: [1,20,30], 1: [11,23,45] }
0
Amit Yadav 6 oct. 2019 a las 12:23
58255798