使用numpy在2D和1D数组之间进行逐元素isin比较

Question

我有一个相当简单的情景，我想测试一个二维数组的两个元素（分别）是否是更大数组的成员 - 例如：

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中的元素是否存在于两个子数组之一。相反，我想要一个二维数组，显示每个子数组中的两个元素的情况 - 所以：

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

希望这有意义，非常感谢您的帮助。

- Chris J Harris

Answer 1

基于广播的广告

一个简单的方法是使用 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不一定是排序的，我们需要使用searchsorted函数的sorter参数-

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

- Divakar

---

你能稍微详细解释一下广播解决方案是如何工作的吗？ - uller