Python 在两个 n x 1 的 numpy 数组中查找匹配项

Question

我已经找了一段时间想解决这样的问题：

import numpy

a=[numpy.array([1,2]),numpy.array([2,2]),numpy.array([3,2]),numpy.array([4,2])]
b=[numpy.array([2,2]),numpy.array([3,2]),numpy.array([6,2]),numpy.array([5,2]),numpy.array([5,2])]

ya=numpy.array([1,2,3,4])
size_a=len(a)
size_b=len(b)
yb=numpy.empty((size_b,1))
yb.fill(numpy.nan)

for i in xrange(size_b):
    for j in xrange(size_a):
        if numpy.array_equiv(yb,ya):
            ya[i]=yb[j]

我只是想用数组b中某个元素匹配到的索引位置的ya值来填充yb。因为yb的长度比ya长，所以在循环结束时，yb的后面部分出现“nan”是正常的。下面的代码运行时间太长了。实际上，我不知道它是否有效，因为我没有等到循环结束就停止了...

在实际情况下，ya和yb的长度分别是7007和3525。

有没有其他方法可以实现我的目标呢？

Answer 1

要在数组列表之间找到匹配项，最简单的方法就是把这些列表调整成相同的形状，比如 n x m 这样的格式。你可以用 np.tile 来做到这一点，但使用 stride_tricks 会更快一些：

a = np.array(a)
b = np.array(b)
shape = (2, a.shape[0], b.shape[0])
from numpy.lib.stride_tricks import as_strided
a = as_strided(a, shape=shape, strides=(a.strides[1], a.strides[0], 0))
b = as_strided(b, shape=shape, strides=(b.strides[1], 0, b.strides[0]))
np.where(np.all(a == b, axis=0))

这样就得到了结果

(array([1, 2]), array([0, 1]))

也就是说，a[1] == b[0]，还有 a[2] == b[1]，没有其他的匹配。