Python 在3d numpy数组中查找特定行的索引

我正在寻找一种在numpy中查找3d数组中外部特定行的索引的方法。一个例子是找到RBG图像中给定颜色集的所有出现点，并获取像素坐标

这表明中的

操作符可以在数组中表现得很奇怪，这更接近于二维数组

假设我们有一个三维数组
Z
，它的尺寸
（x，y，Z）
，并且
[s0，s1]
我们想要匹配的第三维行

Z = np.zeros((10,20,3), dtype=int)
s0 = np.array([1,2,3])
s1 = np.array([4,5,6])
Z[1,2] = s0
Z[4,5] = s1

我想要所有
（x，y）
，其中
z
等于
s0
或
s1

到目前为止,，
argwhere
返回
s0
中的一个元素位于
Z
中的每个匹配：

> np.argwhere(s0 == Z) 
array([[1, 2, 0],
       [1, 2, 1],
       [1, 2, 2]])

in1d
返回一个布尔1D数组，其中s0或s1中的元素匹配为True：

> np.in1d(Z, [s0,s1])

如果我试着用这种方式：

> Zravel = np.ascontiguousarray(a).view([('', a.dtype)] * a.shape[-1]).ravel()
> np.all(np.in1d(Zravel, [s0, s1]) == False)

所有元素均为
False

有什么想法吗？
会使它的输入变得平坦。因此，您可以为它提供
Z
和堆叠版本的
s0
，
s1
，从而为我们提供一个布尔数组，该数组可以被重塑为与
Z
形状相同的数组。然后，您需要检查其中是否有匹配索引的所有真实行。实现如下所示-

S = np.row_stack((s0,s1))
out = np.where((np.in1d(Z,S).reshape(Z.shape)).all(2))

你也可以用这样的方法来解决它-

out = np.where(((Z == S[:,None,None,:]).all(3)).any(0))

S = np.concatenate((s0,s1)).reshape(-1,s0.size)

如果希望输出堆叠在阵列中-

outarr = np.column_stack((out))

要创建
S
，可以将
np.row\u stack
替换为
np.concatenate
，这样可能会更快-

out = np.where(((Z == S[:,None,None,:]).all(3)).any(0))

S = np.concatenate((s0,s1)).reshape(-1,s0.size)

样本运行-

In [145]: Z = np.zeros((10,20,3), dtype=int)
     ...: s0 = np.array([1,2,3])
     ...: s1 = np.array([4,5,6])
     ...: Z[1,2] = s0
     ...: Z[4,5] = s1
     ...: 

In [146]: np.where(((Z == S[:,None,None,:]).all(3)).any(0))
Out[146]: (array([1, 4]), array([2, 5]))

In [147]: np.where((np.in1d(Z,S).reshape(Z.shape)).all(2))
Out[147]: (array([1, 4]), array([2, 5]))

In [148]: np.column_stack((np.where(((Z == S[:,None,None,:]).all(3)).any(0))))
Out[148]: 
array([[1, 2],
       [4, 5]])

In [149]: np.column_stack((np.where((np.in1d(Z,S).reshape(Z.shape)).all(2))))
Out[149]: 
array([[1, 2],
       [4, 5]])

所以，诀窍是将S对象膨胀到与Z相匹配的尺寸以进行广播？@toine Yup！那么有多少这样的s0，s1。。你愿意吗？如果只有两个：s0和s1，你根本不需要广播！我只是假设了一个要解决的一般情况，即S将存储s0、s1、s2、s3等等。我将有更多，正如您所猜测的：）这样的解决方案看起来不错。@toine真棒！所以只需在第一行添加这些。此外，您还可以将
np.row\u stack
替换为
np.concatenate
，以提高性能。