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]
3 respuestas
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]
¿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.
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] }