find array elements that are members of set array in python ()

问题

I am new to python and my question might seem very badly explained because I have background in MATLAB. usually in MATLAB if we have 1000 arrays of 15*15 we define a cell or a 3D matrix which each element is a matrix with size (15*15).

Now in python: (using numpy library) I have a ndarray A with shape (1000,15,15). I have another ndarry B with shape (500,15,15).

I am trying to find elements in A that are member in B too. I am specifically looking for a vector to be returned with the indices of elements in A that are found in B too.

Usually in MATLAB I reshape them to make 2D arrays (1000*225) and (500*225) and use 'ismember' function, passing 'rows' argument to find and return the index of the similar rows.

Are there any similar functions in numpy (or any other library) to do the same? I am trying to aviod for loop.

Thanks

回答1:

Here'e one approach using views largely based on this post -

# 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))

Sample run -

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])

来源：https://stackoverflow.com/questions/43480797/find-array-elements-that-are-members-of-set-array-in-python