在python()中查找作为set数组成员的数组元素

时间:2017-04-18 19:39:06

标签: python arrays numpy find

我是python的新手,我的问题可能看起来很糟糕,因为我有MATLAB的背景知识。 通常在MATLAB中,如果我们有1000个15 * 15的数组,我们定义一个单元格或一个3D矩阵,每个元素是一个大小为(15 * 15)的矩阵。

现在在python中:(使用numpy库) 我有一个有形状的ndarray A(1000,15,15)。 我有另一个有形状的ndarry B(500,15,15)。

我正在尝试在A中找到也是B成员的元素。 我特意寻找一个带有A中元素索引的向量返回,也可以在B中找到。

通常在MATLAB中我重塑它们以制作2D阵列(1000 * 225)和(500 * 225)并使用' ismember'功能,传递'行'查找和返回相似行的索引的参数。

在numpy(或任何其他库)中是否有类似的功能可以做同样的事情? 我想尝试循环。

由于

1 个答案:

答案 0 :(得分:1)

这是一种主要基于this post -

使用views的方法
# Based on https://stackoverflow.com/a/41417343/3293881 by @Eric
def get_index_matching_elems(a, b):
    # check that casting to void will create equal size elements
    assert a.shape[1:] == b.shape[1:]
    assert a.dtype == b.dtype

    # compute dtypes
    void_dt = np.dtype((np.void, a.dtype.itemsize * np.prod(a.shape[1:])))

    # convert to 1d void arrays
    a = np.ascontiguousarray(a)
    b = np.ascontiguousarray(b)
    a_void = a.reshape(a.shape[0], -1).view(void_dt)
    b_void = b.reshape(b.shape[0], -1).view(void_dt)

    # Get indices in a that are also in b
    return np.flatnonzero(np.in1d(a_void, b_void))

示例运行 -

In [87]: # Generate a random array, a
    ...: a = np.random.randint(11,99,(8,3,4))
    ...: 
    ...: # Generate random array, b and set few of them same as in a
    ...: b = np.random.randint(11,99,(6,3,4))
    ...: b[[0,2,4]] = a[[3,6,1]]
    ...: 

In [88]: get_index_matching_elems(a,b)
Out[88]: array([1, 3, 6])