Numpy找到匹配列的索引

时间:2014-06-13 12:14:44

标签: python numpy indexing

我有一个大的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]。有兴趣知道是否有一种有效的方式而不是循环。谢谢!

4 个答案:

答案 0 :(得分:4)

你的问题几乎没有一个好的答案:numpy不太适合这类问题,尽管可以做到。要进行子阵列搜索,如果你的dtype不是浮点数,方法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数组中没有重复的列。

如果你的任何一个数组真的很小,就像在你的例子中那样,很难打败像:

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

但是对于较大的数据集,很难超越排序方法:

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

答案 1 :(得分:1)

这是一种方法,因为数组是预先排序的:

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)

答案 2 :(得分:0)

您可以使用基于字符串的比较,例如使用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数组转换为shape=(18,)数组,其中每个元素都包含相应列的两个元素的数据。同样的想法可以应用于数组b,获得shape=(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)

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

答案 3 :(得分:-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