python - 将 numpy.set1d 推广到 nd-arrays

Question

我正在尝试在 nd-arrays 上使用 numpy 的 setdiff1d 函数：

import numpy as np

#a,b being ndarrays

in_a_not_b = np.setdiff1d(a,b)

但它不起作用，因为它在工作 nd-array 元素方面是明智的。

例如，如果：

a = [[1,2,3],[4,5,6]]
b = [[7,2,3],[4,5,6],[7,8,9]]

我希望输出为：

[[1,2,3]]

但这里是：

[1]

有没有一种简单的方法可以将 setdiff1d 推广到 nd-arrays ？

标签： pythonnumpy

Answer 1

您可以将数组相互广播，以使最后一个维度对齐，然后检查最后一个轴的所有元素是否相等，并且第二个到最后一个轴中的至少一个元素是这种情况（数组检查）：

mask = a[:, None, :] == b[None, :, :]
mask = mask.all(axis=-1).any(axis=-1)
result = a[~mask]

元素的顺序在这里很重要，所以如果不应该，那么您需要沿最后一个轴对数组进行排序。