python - B 在 A 中的值(value)位置索引

Question

我在这里遇到了一个棘手的问题......

给定两个数组A和B

A = np.array([8, 5, 3, 7])
B = np.array([5, 5, 7, 8, 3, 3, 3])

我想用 A 中该值的索引替换 B 中的值。在这个示例中，它看起来像:

[1, 1, 3, 0, 2, 2, 2]

对于我正在处理的问题，A 和 B 包含相同的一组值，并且 A 中的所有条目都是唯一的。

解决这个问题的简单方法是使用类似的东西:

for idx in range(len(A)):
    ind = np.where(B == A[idx])[0]
    B_new[ind] = A[idx]

但是我正在使用的 B 数组包含将近一百万个元素，并且使用 for 循环变得非常慢。必须有一种方法可以对此进行矢量化，但我无法弄清楚。我最接近的是做类似的事情

np.intersect1d(A, B, return_indices=True)

但是这只会让我第一次出现 A 的每个元素在 B 中。有什么建议吗？

Answer 1

@mozway 的解决方案适用于小数组但不适用于大数组，因为它在 O(n**2) 时间内运行(即二次时间，请参阅 time complexity 了解更多信息).这是一个基于快速二分搜索的在 O(n log n) 时间(即准线性)运行的大数组的更好的解决方案:

unique_values, index = np.unique(A, return_index=True)
result = index[np.searchsorted(unique_values, B)]