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()

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

如果您的阵列中有一个非常小，如您的示例中所示，则很难击败以下类型的阵列：

^{pr2}$

但对于更大的数据集，排序方法是很难做到的：

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

Answer 2

如果数组是预先排序的，下面是一种方法：

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]
])

def search2D(A, B):
    to_find_and_bounds = zip(
        B[1],
        numpy.searchsorted(A[0], B[0], side="left"),
        numpy.searchsorted(A[0], B[0], side="right")
    ) 

    for to_find, left, right in to_find_and_bounds:
        offset = numpy.searchsorted(A[1, left:right], to_find)
        yield offset + left

list(search2D(A, B))
#>>> [0, 6, 14]

这是O(len B · log len A)。在

对于未排序的数组，可以执行间接排序：

^{pr2}$

如果需要一个索引的多个结果，请尝试

for to_find, left, right in to_find_and_bounds:
    offset_left = numpy.searchsorted(A[1, left:right], to_find, side="left")
    offset_right = numpy.searchsorted(A[1, left:right], to_find, side="right")
    yield from range(offset_left + left, offset_right + left)

Answer 3

您可以使用np.char.array进行基于字符串的比较

ca = np.char.array(a)[0,:] + np.char.array(a)[1,:]
cb = np.char.array(b)[0,:] + np.char.array(b)[1,:]
np.where(np.in1d(ca, cb))[0]
#array([ 0,  6, 14], dtype=int64)

---

编辑：

您还可以操作数组dtype，以便将a数组转换为具有{}的数组，其中每个元素包含相应列的两个元素的数据。同样的思想也可以应用到数组b，得到{}。然后使用np.where(np.in1d())获得索引：

^{pr2}$

这个想法类似于基于字符串的方法。在