Question

我有一个numpy ndarray，其中第一列是用户ID，第二列是某种产品ID。获取给定用户ID的所有产品ID的最快方法是什么？

我也一直在阅读numpy文档和本手册（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

RestoredData_Array = np.dot(featuresT.reshape(2,2), FinalData1_Matrix.reshape(2,20))
RestoredData_Array[0],RestoredData_Array[1] = RestoredData_Array[1],RestoredData_Array[0]
print(RestoredData_Array)

尝试一下。我提供了两种执行同一操作的方法。使用您喜欢的任何一种。

复杂度在O（n）附近，应该足够好并且足够快。

如果您打算多次使用此函数，我建议编写另一个函数，该函数将返回一个字典，其中user_id为键，product_ids为值。与每次迭代整个列表相比，这将使整体操作真正有效。

使用示例输入，该字典将如下所示：

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)

Answer 2

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-liner-就执行速度而言，如果您不使用numpy数组，则Amit Yadav提出的纯Python解决方案是最有效的可能会更快。

Answer 3

这可以通过一种简单的方式实现

test = np.array(test)

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

从给定的第一列值的ndarray提取所有第二列值

3 个答案: