如何在另一个矩阵中选择一个矩阵的向量

时间:2017-05-22 00:47:22

标签: python pandas numpy

我有两个矩阵:

a = [[1,3,4],[2,5,3],[2,4,6],[6,5,3]] 
b = [[2,4,5],[2,4,6],[1,3,4]] 

我想在b中选择[2,4,6],[1,3,4],这是在a。

由于a和b很大,

for v in b:
    if  v in a:

很贵。

有人可以告诉我最好的方法吗?

1 个答案:

答案 0 :(得分:1)

对于二维矩阵,您想要的是numpy.in1d的等价物。我刚才写了这样一个函数

def in2d(arr1, arr2):
    """Generalisation of numpy.in1d to 2D arrays"""

    assert arr1.dtype == arr2.dtype

    arr1_view = np.ascontiguousarray(arr1).view(np.dtype((np.void,
         arr1.dtype.itemsize * arr1.shape[1])))
    arr2_view = np.ascontiguousarray(arr2).view(np.dtype((np.void,
         arr2.dtype.itemsize * arr2.shape[1])))
    intersected = np.in1d(arr1_view, arr2_view)
    return intersected.view(np.bool).reshape(-1)

可以找到有关其工作原理的说明here。 你可以使用像这样的功能

In [56]: a = np.array([[1,3,4],[2,5,3],[2,4,6],[6,5,3]])
In [57]: b = np.array([[2,4,5],[2,4,6],[1,3,4]])

In [58]: in2d(b,a)
Out[58]: array([False,  True,  True], dtype=bool) 

它返回一个布尔数组,其中b的元素位于a中。反之亦然

In [59]: in2d(a,b)
Out[59]: array([ True, False,  True, False], dtype=bool)

使用此布尔数组索引a可以提供您想要的内容

In [60]: a[in2d(a,b),:]
Out[60]: 
array([[1, 3, 4],
       [2, 4, 6]])

请注意,您的解决方案(在下面发布)不正确,并且没有按照您的想法执行,因为if v in a搜索所有嵌套数组/列表元素。所以下面的比较是不公平的,不过考虑

def for_loop_and_compare(a,b):
    return np.array([v for v in b if v in a])

和时间

In [61]: a=np.random.randint(0,100,(10000,3))
In [62]: b=np.random.randint(0,100,(1000,3))
In [63]: %timeit for_loop_and_compare(a,b)
10 loops, best of 3: 79 ms per loop
In [64]: %timeit in2d(a,b)
100 loops, best of 3: 3.7 ms per loop