python获取set中某些元素_在python中找到属于set数组成员的数组元素

  • Post author:
  • Post category: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

解决方案

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



版权声明:本文为weixin_39685578原创文章,遵循 CC 4.0 BY-SA 版权协议,转载请附上原文出处链接和本声明。