Question

我有一个numpy数组A，其中包含可以按任意顺序排列的唯一ID，例如A = [1、3、2]。我有第二个numpy数组B，它记录了使用ID的时间-例如B = [3、3、1、3、2、1、2、3、1、1、2、3、3、1]。数组B总是比数组A长得多。

每次在B中使用ID时，我都需要在A中找到ID的索引位置。因此在上面的示例中，我返回的结果将是：result = [1、1、0、1、2、0 ，2、1、0、0、2、1、1、0]。

我已经编写了一个简单的解决方案，该解决方案使用for循环将结果附加到新列表并使用numpy.where来获取正确的结果，但是我无法弄清楚将其向量化的正确语法。

import numpy as np
A = np.array([1, 3, 2])
B = np.array([3, 3, 1, 3, 2, 1, 2, 3, 1, 1, 2, 3, 3, 1])

IdIndxs = []
for ID in B:
    IdIndxs.append(np.where(A == ID)[0][0])

IdIndxs = np.array(IdIndxs)

有人能想出一个简单的基于矢量的解决方案，该解决方案可以快速运行吗？-对于一个典型的问题，其中A的大小为10K-100K，而B的倍数通常为5，因此for循环会变得非常慢。比A大10倍。

我确定解决方案很简单，但是今天我看不到。

Answer 1

numpy-indexed库（免责声明：我是它的作者）旨在提供这些类型的矢量化操作，但出于某些原因numpy却不提供。坦白地说，考虑到矢量化list.index等效项的实用性，它绝对应该以numpy表示；但是numpy是一个发展缓慢的项目，非常重视向后兼容性，而且我认为直到numpy2.0之前我们都不会看到它。但在此之前，可以轻松安装pip和conda。

import numpy_indexed as npi
idx = npi.indices(A, B)

Answer 2

您可以使用此：

import numpy as np

# test data
A = np.array([1, 3, 2])
B = np.array([3, 3, 1, 3, 2, 1, 2, 3, 1, 1, 2, 3, 3, 1])

# get indexes
sorted_keys = np.argsort(A)
indexes = sorted_keys[np.searchsorted(A, B, sorter=sorted_keys)]

输出：

[1 1 0 1 2 0 2 1 0 0 2 1 1 0]

Answer 3

重新设计您的逻辑，但使用列表理解和numpy.fromiter可以提高性能。

IdIndxs = np.fromiter([np.where(A == i)[0][0] for i in B], B.dtype)

关于效果

我已经进行了一项快速测试，将fromiter与您的解决方案进行了比较，但没有看到这种性能提升。即使使用由数百万个元素组成的B数组，它们的顺序也相同。

从另一个未排序的numpy数组中的数据中查找未排序的numpy数组中的值的索引位置

3 个答案:

关于效果