Question

我正在尝试从numpy ndarray中为我的协作过滤项目选择一个子集。

我的阵列有这样的形状：

ratings = np.array(
    [
        (1, 2, 3.0),
        (2, 2, 3.0),
        (4, 1, 2.0),
        (1, 2, 1.0),
    ],
    dtype=[
        ('user_id', np.uint32),
        ('item_id', np.uint32),
        ('score', np.float32)
    ]
)

现在，我想选择user_id属于数组的评级子集。类似于SQL的'WHERE IN'功能。

我能够使用np.in1d：

实现这一目标

subset_of_users = [1, 2]
ratings[np.in1d(ratings['user_id'], subset_of_users)]

我的分析总是表明使用in1d的块是最慢的，它让我觉得可能有更快的替代方法。

非常感谢你的时间。

Answer 1

您可以尝试使用numpy_indexed包（免责声明：我是其作者）。

受您的问题的启发，我已经更改了API，以便我们可以从重复查询中受益：

import numpy_indexed as npi
npi.contains(subset_of_users, ratings['user_id'])

应该从左到右阅读; 'subset包含user_id的元素'，并返回子集中存在的'user_id'的索引。

然而，计算中最昂贵的部分是为用户ID集构建“索引”，这可以预先计算：

index = npi.as_index(ratings['user_id'])
npi.contains(subset_of_users, index)
npi.contains(some_other_subset_of_users, index)

我期望在每个查询的基础上加快速度。

我还整合了一个npi.in_函数，受到Divakar的回答的启发，它允许你编写npi.in_（ratings ['user_id']，subset_of_users），它再次从左向右读; user_id的元素，它们存在于子集'中。但是我认为它比使用contains有点效率低。但这都是猜想;很高兴看到对实际数据的一些比较！

Answer 2

如果您的最大user_id不是太大，您可以使用查找表：

mask_of_users = np.zeros(shape=max(ratings['user_id'])+1, dtype=bool)
mask_of_users[subset_of_users] = True
selected_ratings = ratings[mask_of_users[ratings['user_id']]]

Answer 3

似乎存在np.in1d的瓶颈，所以让我们尝试加速那部分。现在，根据我迄今为止的NumPy经验，我遇到了一个np.searchsorted的替代方案，当第二个数组是一个已排序且唯一的数组或列表时，它可用作np.in1d的替代。与第一个数组中的元素匹配。下面列出的是实施 -

def in1d_replacement(A,B):
    """ in1d replacement using searchsorted with optional 'left', 'right' args.
    """
    # Get left and right sorted indices for A in B
    idx_l = np.searchsorted(B,A,'left')
    idx_r = np.searchsorted(B,A,'right')

    # Look for switching indices between the left and right ones 
    # indicating the matches
    return idx_l != idx_r

运行时测试 -

In [195]: # Random arrays of decent sizes
     ...: nA = 10000
     ...: nB = 1000
     ...: max_num = 100000
     ...: A = np.random.randint(0,max_num,(nA))
     ...: B = np.unique(np.random.randint(0,max_num,(nB)))
     ...: 

In [196]: np.allclose(np.in1d(A,B),in1d_replacement(A,B))
Out[196]: True

In [197]: %timeit np.in1d(A,B)
100 loops, best of 3: 2.73 ms per loop

In [198]: %timeit in1d_replacement(A,B)
100 loops, best of 3: 2.14 ms per loop

所以，如果不是一个巨大的性能，那么会有一些性能提升。

Numpy WHERE IN等价物

3 个答案: