python - Numpy 查找匹配列的索引

Question

我有一个大的 2xn 数组 A 和一个较小的 2xn 数组 B。B 中的所有列都可以在 A 中找到。我希望通过匹配 B 中的列来查找 A 的索引。例如，

import numpy

A = numpy.array([
    [101, 101, 101, 102, 102, 103, 103, 104, 105, 106, 107, 108, 108, 109, 109, 110, 110, 211],
    [102, 103, 105, 104, 106, 109, 224, 109, 110, 110, 108, 109, 110, 211, 212, 211, 212, 213]
])

B = numpy.array([
    [101, 103, 109],
    [102, 224, 212]
])

我正在寻找的答案是[0,6,14]。有兴趣知道是否有一种有效的方法而不是循环。谢谢!

Answer 1

对于您的问题几乎没有一个好的答案:numpy 不太适合此类问题，尽管可以做到。要进行子数组搜索，如果您的数据类型不是浮点，则方法 here可能是你最好的选择。您可以从以下内容开始:

AA = np.ascontiguousarray(A.T)
BB = np.ascontiguousarray(B.T)

dt = np.dtype((np.void, AA.dtype.itemsize * AA.shape[1]))
AA = AA.view(dt).ravel()
BB = BB.view(dt).ravel()

现在只是在另一个一维数组中搜索一个一维数组中的项目，这非常简单，假设原始 A 数组中没有重复的列。

如果您的任何一个数组非常小(如您的示例所示)，则很难击败以下内容:

indices = np.argmax(AA == BB[:, None], axis = 1)

但对于较大的数据集，很难击败排序方法:

sorter = np.argsort(AA)
sorted_indices = np.searchsorted(AA, BB, sorter=sorter)
indices = sorter[sorted_indices]