python - 等同于Numpy

我正在尝试从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的块是最慢的，这让我想到也许有更快的替代方法。

非常感谢您的宝贵时间。

最佳答案

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

受您的问题启发，我对API做了一些更改，以使我们可以从重复的查询中受益：

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

应该从左到右阅读； “子集包含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)

我希望每个查询的速度会更快一些。

Ive还结合了Divakar的答案，集成了npi.in_函数，该函数使您可以编写npi.in_（ratings ['user_id']，subset_of_users），该函数再次从左至右读取；子集'中存在的user_id元素。但是我希望它比使用contains效率低一些。但这全是猜测。很高兴看到实际数据的一些比较！