从给定第一列值的数据数组中提取所有第二列值

Question

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

我也读过numpy doc和这本手册（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])

…尽管这是一个很好的单行程序—关于执行速度，如果您不想使用numpy数组，那么像Amit Yadav提出的纯Python解决方案很可能更快

Answer 2

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\u id作为键，product\u id作为值。这将使整个操作比每次遍历完整列表更有效

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

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

Answer 3

这可以通过这样一种简单的方法来实现

test = np.array(test)

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