按元素使用numpy .isin

时间:2018-12-05 11:39:59

标签: python-3.x numpy numpy-broadcasting

我有一个非常简单的场景,我想测试二维数组的两个元素是否(分别)是较大数组的成员-例如:

full_array = np.array(['A','B','C','D','E','F'])
sub_arrays = np.array([['A','C','F'],
                       ['B','C','E']])
np.isin(full_array, sub_arrays)

这给了我一维输出:

array([ True,  True,  True, False,  True,  True])

显示full_array的元素是否存在于两个子数组中的任何一个中。我想要一个二维数组,为sub_arrays中的两个元素中的每个元素显示相同的内容-所以:

array([[ True,  False,  True, False,  False,  True],
       [ False, True,   True, False,  True,  False]])

希望如此,希望得到的任何帮助。

1 个答案:

答案 0 :(得分:4)

基于广播的

在展开其中一个数组然后沿各自的轴进行任意归约后,可以使用broadcasting进行简单的处理-

In [140]: (full_array==sub_arrays[...,None]).any(axis=1)
Out[140]: 
array([[ True, False,  True, False, False,  True],
       [False,  True,  True, False,  True, False]])

使用searchsorted

特定案例#1

full_array进行排序,并且sub_arrays中的所有元素至少出现在full_array中的某个位置,我们也可以使用np.searchsorted-

idx = np.searchsorted(full_array, sub_arrays)
out = np.zeros((sub_arrays.shape[0],len(full_array)),dtype=bool)
np.put_along_axis(out, idx, 1, axis=1)

特定案例#2

在对full_array进行排序的情况下,如果不能保证sub_arrays中的所有元素都至少存在于full_array中的某个地方,我们需要执行一个额外的步骤-

idx = np.searchsorted(full_array, sub_arrays)
idx[idx==len(full_array)] = 0
out = np.zeros((sub_arrays.shape[0],len(full_array)),dtype=bool)
np.put_along_axis(out, idx, full_array[idx] == sub_arrays, axis=1)

一般情况

对于full_array的真正通用情况(不一定要排序),我们需要将sorter arg与searchsorted-

def isin2D(full_array, sub_arrays):
    out = np.zeros((sub_arrays.shape[0],len(full_array)),dtype=bool)
    sidx = full_array.argsort()
    idx = np.searchsorted(full_array, sub_arrays, sorter=sidx)
    idx[idx==len(full_array)] = 0
    idx0 = sidx[idx]
    np.put_along_axis(out, idx0, full_array[idx0] == sub_arrays, axis=1)
    return out

样品运行-

In [214]: full_array
Out[214]: array(['E', 'F', 'A', 'B', 'D', 'C'], dtype='|S1')

In [215]: sub_arrays
Out[215]: 
array([['Z', 'C', 'F'],
       ['B', 'C', 'E']], dtype='|S1')

In [216]: isin2D(full_array, sub_arrays)
Out[216]: 
array([[False,  True, False, False, False,  True],
       [ True, False, False,  True, False,  True]])
相关问题