我有两个矩阵:
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:
很贵。
有人可以告诉我最好的方法吗?
答案 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