在Python中查找集合数组的成员元素

Question

我刚开始学习Python，我的问题可能解释得很差，因为我有MATLAB的背景。通常，在MATLAB中，如果我们有1000个15*15的数组，我们定义一个单元格或3D矩阵，其中每个元素都是大小为(15*15)的矩阵。

现在在Python中（使用numpy库），我有一个形状为(1000, 15, 15)的ndarray A和另一个形状为(500, 15, 15)的ndarry B。

我试图找到A中也属于B的元素。我特别要求返回一个向量，其中包含A中与B相同元素的索引。

通常在MATLAB中，我将它们重塑为2D数组（1000*225）和（500*225），并使用'ismember'函数传递'rows'参数来查找并返回相似行的索引。

在numpy（或任何其他库）中是否有类似的功能来完成同样的操作呢？我正在尝试避免for循环。

谢谢

- Jahandar Jahanipour

---

看看这个是否有帮助：this。 - Divakar

实际上这返回相似的矩阵。我正在寻找它们在原始矩阵中的索引。 - Jahandar Jahanipour

这个回答解决了你的问题吗？Python中类似MATLAB的“ismember”函数的等价实现 - Cris Luengo

Answer 1

这里提供一种基于视图（views）的方法，主要参考这篇文章 -

# Based on https://dev59.com/1Z3ha4cB1Zd3GeqPZ853#41417343 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])

- Divakar