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库里有你需要的一切。我假设这些数组是没有排序的，当然你可以根据自己的需要来改进下面的代码：

import numpy as np

a = np.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 = np.array([[101, 103, 109],
     [102, 224, 212]])

idxs = []

for i in range(np.shape(b)[1]):
    for j in range(np.shape(a)[1]):
        if np.array_equal(b[:,i],a[:,j]):
            idxs.append(j)

print idxs

Answer 2

你可以使用基于字符串的比较方法，比如用 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 数组变成一个形状为 (18,) 的数组，每个元素包含对应列的两个元素的数据。这个思路同样可以用在数组 b 上，得到的形状是 (3,)。接着，你可以用 np.where(np.in1d()) 来获取索引：

nrows = a.shape[0]
ta = np.ascontiguousarray(a.T).view(np.dtype((np.void, a.itemsize*nrows))).flatten()
tb = np.ascontiguousarray(b.T).view(np.dtype((np.void, b.itemsize*nrows))).flatten()
np.where(np.in1d(ta, tb))[0]
#array([ 0,  6, 14], dtype=int64)

这个思路和基于字符串的方法是类似的。

Answer 3

对于你的问题，几乎没有一个好的答案：numpy并不太适合这种类型的问题，虽然是可以解决的。要进行子数组搜索，如果你的数据类型不是浮点数，那么这里的方法可能是你最好的选择。你可以从这样的代码开始：

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]

Answer 4

这里有一种方法，前提是数组已经排好序：

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)，意思是处理的速度和数组的长度有关。

如果数组没有排好序，你可以先进行间接排序：

sorter = numpy.lexsort(A[::-1])
sorted_copy = A.T[sorter].T

sorter[list(search2D(sorted_copy, B))]
#>>> array([ 3,  6, 14])

如果你需要从一个索引得到多个结果，可以试试：

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)